diff --git a/MODEL_ZOO.md b/MODEL_ZOO.md index 0436dae72..191ed594e 100644 --- a/MODEL_ZOO.md +++ b/MODEL_ZOO.md @@ -12,6 +12,7 @@ | SlowFast | R50 | 3 x 10 | 4 x 16 | 75.6 | 92.0 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/kinetics400/SLOWFAST_4x16_R50.pkl) | Kinetics/c2/SLOWFAST_4x16_R50 | K400 | | SlowFast | R50 | 3 x 10 | 8 x 8 | 77.0 | 92.6 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/kinetics400/SLOWFAST_8x8_R50.pkl) | Kinetics/c2/SLOWFAST_8x8_R50 | K400 | | MViTv1 | B-Conv | 1 x 5 | 16 x 4 | 78.4 | 93.5 | [`link`](https://drive.google.com/file/d/194gJinVejq6A1FmySNKQ8vAN5-FOY-QL/view?usp=sharing) | Kinetics/MVIT_B_16x4_CONV | K400 | +| rev-MViT | B-Conv | 1 x 5 | 16 x 4 | 78.4 | 93.4 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_MVIT_B_16x4.pyth) | Kinetics/REV_MVIT_B_16x4_CONV | K400 | | MViTv1 | B-Conv | 1 x 5 | 32 x 3 | 80.4 | 94.8 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvit/k400.pyth) | Kinetics/MVIT_B_32x3_CONV | K400 | | MViTv1 | B-Conv | 1 x 5 | 32 x 3 | 83.9 | 96.5 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvit/k600.pyth) | Kinetics/MVIT_B_32x3_CONV_K600 | K600 | | MViTv2 | S | 1 x 5 | 16 x 4 | 81.0 | 94.6 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvitv2/pysf_video_models/MViTv2_S_16x4_k400_f302660347.pyth) | Kinetics/MVITv2_S_16x4 | K400 | @@ -68,10 +69,13 @@ We also release the imagenet pretrained model if finetuning from ImageNet is pre | architecture | size | Top1 | Top5 | model | Config | | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | -| ResNet | R50 | 23.6 | 6.8 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/kinetics400/R50_IN1K.pyth) | ImageNet/RES_R50 | -| MVIT | B-16-Conv | 17.1 | 3.7 | [`link`](https://drive.google.com/file/d/1dYYqUB-3DSgBVc9d6o-rW8ojtVsrFLgp/view?usp=sharing) | ImageNet/MVIT_B_16_CONV | - +| ResNet | R50 | 76.4 | 93.2 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/kinetics400/R50_IN1K.pyth) | ImageNet/RES_R50 | +| MVIT | B-16-Conv | 82.9 | 96.3 | [`link`](https://drive.google.com/file/d/1dYYqUB-3DSgBVc9d6o-rW8ojtVsrFLgp/view?usp=sharing) | ImageNet/MVIT_B_16_CONV | +| rev-VIT | Small | 79.9 | 94.9 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_VIT_S.pyth) | ImageNet/REV_VIT_S.yaml | +| rev-VIT | Base | 81.8 | 95.6 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_VIT_B.pyth) | ImageNet/REV_VIT_B.yaml | +| rev-MVIT | Base | 82.9* | 96.3 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_MVIT_B.pyth) | ImageNet/REV_MVIT_B_16_CONV.yaml | +*please refer to [Reversible Model Zoo](projects/rev/README.md). ## PyTorchVideo diff --git a/README.md b/README.md index fcd49dbde..25093efad 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,12 @@ PySlowFast is an open source video understanding codebase from FAIR that provide - [Non-local Neural Networks](https://arxiv.org/abs/1711.07971) - [A Multigrid Method for Efficiently Training Video Models](https://arxiv.org/abs/1912.00998) - [X3D: Progressive Network Expansion for Efficient Video Recognition](https://arxiv.org/abs/2004.04730) -- [Multiscale Vision Transformers](https://arxiv.org/abs/2104.11227.pdf) +- [Multiscale Vision Transformers](https://arxiv.org/abs/2104.11227) +- [A Large-Scale Study on Unsupervised Spatiotemporal Representation Learning](https://arxiv.org/abs/2104.14558) +- [MViTv2: Improved Multiscale Vision Transformers for Classification and Detection](https://arxiv.org/abs/2112.01526) +- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://arxiv.org/abs/2112.09133) +- [Masked Autoencoders As Spatiotemporal Learners](https://arxiv.org/abs/2205.09113) +- [Reversible Vision Transformers](https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf)
@@ -23,9 +28,14 @@ The goal of PySlowFast is to provide a high-performance, light-weight pytorch co - Non-local Network - X3D - MViTv1 and MViTv2 +- Rev-ViT and Rev-MViT ## Updates + - We now [Reversible Vision Transformers](https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf). Both Reversible ViT and MViT models released. See [`projects/rev`](./projects/rev/README.md). + - We now support [MAE for Video](https://arxiv.org/abs/2104.11227.pdf). See [`projects/mae`](./projects/mae/README.md) for more information. + - We now support [MaskFeat](https://arxiv.org/abs/2112.09133). See [`projects/maskfeat`](./projects/maskfeat/README.md) for more information. - We now support [MViTv2](https://arxiv.org/abs/2104.11227.pdf) in PySlowFast. See [`projects/mvitv2`](./projects/mvitv2/README.md) for more information. + - We now support [A Large-Scale Study on Unsupervised Spatiotemporal Representation Learning](https://arxiv.org/abs/2104.14558). See [`projects/contrastive_ssl`](./projects/contrastive_ssl/README.md) for more information. - We now support [Multiscale Vision Transformers](https://arxiv.org/abs/2104.11227.pdf) on Kinetics and ImageNet. See [`projects/mvit`](./projects/mvit/README.md) for more information. - We now support [PyTorchVideo](https://github.com/facebookresearch/pytorchvideo) models and datasets. See [`projects/pytorchvideo`](./projects/pytorchvideo/README.md) for more information. - We now support [X3D Models](https://arxiv.org/abs/2004.04730). See [`projects/x3d`](./projects/x3d/README.md) for more information. diff --git a/slowfast/utils/ava_evaluation/README.md b/ava_evaluation/README.md similarity index 100% rename from slowfast/utils/ava_evaluation/README.md rename to ava_evaluation/README.md diff --git a/slowfast/utils/ava_evaluation/ava_action_list_v2.1_for_activitynet_2018.pbtxt.txt b/ava_evaluation/ava_action_list_v2.1_for_activitynet_2018.pbtxt.txt similarity index 100% rename from slowfast/utils/ava_evaluation/ava_action_list_v2.1_for_activitynet_2018.pbtxt.txt rename to ava_evaluation/ava_action_list_v2.1_for_activitynet_2018.pbtxt.txt diff --git a/slowfast/utils/ava_evaluation/label_map_util.py b/ava_evaluation/label_map_util.py similarity index 94% rename from slowfast/utils/ava_evaluation/label_map_util.py rename to ava_evaluation/label_map_util.py index ce8933dd2..f78c00e48 100644 --- a/slowfast/utils/ava_evaluation/label_map_util.py +++ b/ava_evaluation/label_map_util.py @@ -14,12 +14,8 @@ # ============================================================================== """Label map utility functions.""" -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) +from __future__ import absolute_import, division, print_function, unicode_literals + import logging # from google.protobuf import text_format @@ -71,9 +67,7 @@ def get_max_label_map_index(label_map): return max([item.id for item in label_map.item]) -def convert_label_map_to_categories( - label_map, max_num_classes, use_display_name=True -): +def convert_label_map_to_categories(label_map, max_num_classes, use_display_name=True): """Loads label map proto and returns categories list compatible with eval. This function loads a label map and returns a list of dicts, each of which @@ -111,8 +105,7 @@ def convert_label_map_to_categories( for item in label_map.item: if not 0 < item.id <= max_num_classes: logging.info( - "Ignore item %d since it falls outside of requested " - "label range.", + "Ignore item %d since it falls outside of requested label range.", item.id, ) continue @@ -136,10 +129,10 @@ def load_labelmap(path): """ with open(path, "r") as fid: label_map_string = fid.read() - label_map = string_int_label_map_pb2.StringIntLabelMap() + label_map = string_int_label_map_pb2.StringIntLabelMap() # noqa try: - text_format.Merge(label_map_string, label_map) - except text_format.ParseError: + text_format.Merge(label_map_string, label_map) # noqa + except text_format.ParseError: # noqa label_map.ParseFromString(label_map_string) _validate_label_map(label_map) return label_map diff --git a/slowfast/utils/ava_evaluation/metrics.py b/ava_evaluation/metrics.py similarity index 92% rename from slowfast/utils/ava_evaluation/metrics.py rename to ava_evaluation/metrics.py index c59a68d09..841a177ed 100644 --- a/slowfast/utils/ava_evaluation/metrics.py +++ b/ava_evaluation/metrics.py @@ -14,7 +14,9 @@ # ============================================================================== """Functions for computing metrics like precision, recall, CorLoc and etc.""" + from __future__ import division + import numpy as np @@ -38,7 +40,7 @@ def compute_precision_recall(scores, labels, num_gt): """ if ( not isinstance(labels, np.ndarray) - or labels.dtype != np.bool + or labels.dtype != bool or len(labels.shape) != 1 ): raise ValueError("labels must be single dimension bool numpy array") @@ -47,9 +49,7 @@ def compute_precision_recall(scores, labels, num_gt): raise ValueError("scores must be single dimension numpy array") if num_gt < np.sum(labels): - raise ValueError( - "Number of true positives must be smaller than num_gt." - ) + raise ValueError("Number of true positives must be smaller than num_gt.") if len(scores) != len(labels): raise ValueError("scores and labels must be of the same size.") @@ -92,13 +92,11 @@ def compute_average_precision(precision, recall): if precision is None: if recall is not None: raise ValueError("If precision is None, recall must also be None") - return np.NAN + return np.nan - if not isinstance(precision, np.ndarray) or not isinstance( - recall, np.ndarray - ): + if not isinstance(precision, np.ndarray) or not isinstance(recall, np.ndarray): raise ValueError("precision and recall must be numpy array") - if precision.dtype != np.float or recall.dtype != np.float: + if precision.dtype != float or recall.dtype != float: raise ValueError("input must be float numpy array.") if len(precision) != len(recall): raise ValueError("precision and recall must be of the same size.") @@ -125,9 +123,7 @@ def compute_average_precision(precision, recall): return average_precision -def compute_cor_loc( - num_gt_imgs_per_class, num_images_correctly_detected_per_class -): +def compute_cor_loc(num_gt_imgs_per_class, num_images_correctly_detected_per_class): """Compute CorLoc according to the definition in the following paper. https://www.robots.ox.ac.uk/~vgg/rg/papers/deselaers-eccv10.pdf diff --git a/slowfast/utils/ava_evaluation/np_box_list.py b/ava_evaluation/np_box_list.py similarity index 95% rename from slowfast/utils/ava_evaluation/np_box_list.py rename to ava_evaluation/np_box_list.py index a3bd38e4b..34dabc4b5 100644 --- a/slowfast/utils/ava_evaluation/np_box_list.py +++ b/ava_evaluation/np_box_list.py @@ -15,16 +15,12 @@ """Numpy BoxList classes and functions.""" -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np -class BoxList(object): +class BoxList: """Box collection. BoxList represents a list of bounding boxes as numpy array, where each @@ -51,9 +47,7 @@ def __init__(self, data): if len(data.shape) != 2 or data.shape[1] != 4: raise ValueError("Invalid dimensions for box data.") if data.dtype != np.float32 and data.dtype != np.float64: - raise ValueError( - "Invalid data type for box data: float is required." - ) + raise ValueError("Invalid data type for box data: float is required.") if not self._is_valid_boxes(data): raise ValueError( "Invalid box data. data must be a numpy array of " diff --git a/slowfast/utils/ava_evaluation/np_box_list_ops.py b/ava_evaluation/np_box_list_ops.py similarity index 95% rename from slowfast/utils/ava_evaluation/np_box_list_ops.py rename to ava_evaluation/np_box_list_ops.py index 7921cff2e..eafa697dd 100644 --- a/slowfast/utils/ava_evaluation/np_box_list_ops.py +++ b/ava_evaluation/np_box_list_ops.py @@ -19,18 +19,15 @@ * Areas: compute bounding box areas * IOU: pairwise intersection-over-union scores """ -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) + +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np from . import np_box_list, np_box_ops -class SortOrder(object): +class SortOrder: """Enum class for sort order. Attributes: @@ -239,9 +236,7 @@ def non_max_suppression( return gather(boxlist, np.array(selected_indices)) -def multi_class_non_max_suppression( - boxlist, score_thresh, iou_thresh, max_output_size -): +def multi_class_non_max_suppression(boxlist, score_thresh, iou_thresh, max_output_size): """Multi-class version of non maximum suppression. This op greedily selects a subset of detection bounding boxes, pruning @@ -284,7 +279,7 @@ def multi_class_non_max_suppression( elif len(scores.shape) == 2: if scores.shape[1] is None: raise ValueError( - "scores field must have statically defined second " "dimension" + "scores field must have statically defined second dimension" ) else: raise ValueError("scores field must be of rank 1 or 2") @@ -334,9 +329,7 @@ def scale(boxlist, y_scale, x_scale): y_max = y_scale * y_max x_min = x_scale * x_min x_max = x_scale * x_max - scaled_boxlist = np_box_list.BoxList( - np.hstack([y_min, x_min, y_max, x_max]) - ) + scaled_boxlist = np_box_list.BoxList(np.hstack([y_min, x_min, y_max, x_max])) fields = boxlist.get_extra_fields() for field in fields: @@ -376,9 +369,9 @@ def clip_to_window(boxlist, window): ) clipped = _copy_extra_fields(clipped, boxlist) areas = area(clipped) - nonzero_area_indices = np.reshape( - np.nonzero(np.greater(areas, 0.0)), [-1] - ).astype(np.int32) + nonzero_area_indices = np.reshape(np.nonzero(np.greater(areas, 0.0)), [-1]).astype( + np.int32 + ) return gather(clipped, nonzero_area_indices) @@ -398,9 +391,7 @@ def prune_non_overlapping_boxes(boxlist1, boxlist2, minoverlap=0.0): A pruned boxlist with size [N', 4]. """ intersection_over_area = ioa(boxlist2, boxlist1) # [M, N] tensor - intersection_over_area = np.amax( - intersection_over_area, axis=0 - ) # [N] tensor + intersection_over_area = np.amax(intersection_over_area, axis=0) # [N] tensor keep_bool = np.greater_equal(intersection_over_area, np.array(minoverlap)) keep_inds = np.nonzero(keep_bool)[0] new_boxlist1 = gather(boxlist1, keep_inds) @@ -472,9 +463,7 @@ def concatenate(boxlists, fields=None): raise ValueError("boxlists should have nonzero length") for boxlist in boxlists: if not isinstance(boxlist, np_box_list.BoxList): - raise ValueError( - "all elements of boxlists should be BoxList objects" - ) + raise ValueError("all elements of boxlists should be BoxList objects") concatenated = np_box_list.BoxList( np.vstack([boxlist.get() for boxlist in boxlists]) ) @@ -527,12 +516,11 @@ def filter_scores_greater_than(boxlist, thresh): raise ValueError("Scores should have rank 1 or 2") if len(scores.shape) == 2 and scores.shape[1] != 1: raise ValueError( - "Scores should have rank 1 or have shape " - "consistent with [None, 1]" + "Scores should have rank 1 or have shape consistent with [None, 1]" ) - high_score_indices = np.reshape( - np.where(np.greater(scores, thresh)), [-1] - ).astype(np.int32) + high_score_indices = np.reshape(np.where(np.greater(scores, thresh)), [-1]).astype( + np.int32 + ) return gather(boxlist, high_score_indices) @@ -580,9 +568,7 @@ def _copy_extra_fields(boxlist_to_copy_to, boxlist_to_copy_from): boxlist_to_copy_to with extra fields. """ for field in boxlist_to_copy_from.get_extra_fields(): - boxlist_to_copy_to.add_field( - field, boxlist_to_copy_from.get_field(field) - ) + boxlist_to_copy_to.add_field(field, boxlist_to_copy_from.get_field(field)) return boxlist_to_copy_to diff --git a/slowfast/utils/ava_evaluation/np_box_mask_list.py b/ava_evaluation/np_box_mask_list.py similarity index 87% rename from slowfast/utils/ava_evaluation/np_box_mask_list.py rename to ava_evaluation/np_box_mask_list.py index 7fd6b341b..c9a462b84 100644 --- a/slowfast/utils/ava_evaluation/np_box_mask_list.py +++ b/ava_evaluation/np_box_mask_list.py @@ -15,12 +15,8 @@ """Numpy BoxMaskList classes and functions.""" -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np from . import np_box_list @@ -55,13 +51,9 @@ def __init__(self, box_data, mask_data): if len(mask_data.shape) != 3: raise ValueError("Invalid dimensions for mask data.") if mask_data.dtype != np.uint8: - raise ValueError( - "Invalid data type for mask data: uint8 is required." - ) + raise ValueError("Invalid data type for mask data: uint8 is required.") if mask_data.shape[0] != box_data.shape[0]: - raise ValueError( - "There should be the same number of boxes and masks." - ) + raise ValueError("There should be the same number of boxes and masks.") self.data["masks"] = mask_data def get_masks(self): diff --git a/slowfast/utils/ava_evaluation/np_box_mask_list_ops.py b/ava_evaluation/np_box_mask_list_ops.py similarity index 93% rename from slowfast/utils/ava_evaluation/np_box_mask_list_ops.py rename to ava_evaluation/np_box_mask_list_ops.py index e91928476..be906023f 100644 --- a/slowfast/utils/ava_evaluation/np_box_mask_list_ops.py +++ b/ava_evaluation/np_box_mask_list_ops.py @@ -19,12 +19,9 @@ * Areas: compute bounding box areas * IOU: pairwise intersection-over-union scores """ -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) + +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np from . import np_box_list_ops, np_box_mask_list, np_mask_ops @@ -91,9 +88,7 @@ def iou(box_mask_list1, box_mask_list2): Returns: a numpy array with shape [N, M] representing pairwise iou scores. """ - return np_mask_ops.iou( - box_mask_list1.get_masks(), box_mask_list2.get_masks() - ) + return np_mask_ops.iou(box_mask_list1.get_masks(), box_mask_list2.get_masks()) def ioa(box_mask_list1, box_mask_list2): @@ -110,9 +105,7 @@ def ioa(box_mask_list1, box_mask_list2): Returns: a numpy array with shape [N, M] representing pairwise ioa scores. """ - return np_mask_ops.ioa( - box_mask_list1.get_masks(), box_mask_list2.get_masks() - ) + return np_mask_ops.ioa(box_mask_list1.get_masks(), box_mask_list2.get_masks()) def gather(box_mask_list, indices, fields=None): @@ -142,15 +135,11 @@ def gather(box_mask_list, indices, fields=None): if "masks" not in fields: fields.append("masks") return box_list_to_box_mask_list( - np_box_list_ops.gather( - boxlist=box_mask_list, indices=indices, fields=fields - ) + np_box_list_ops.gather(boxlist=box_mask_list, indices=indices, fields=fields) ) -def sort_by_field( - box_mask_list, field, order=np_box_list_ops.SortOrder.DESCEND -): +def sort_by_field(box_mask_list, field, order=np_box_list_ops.SortOrder.DESCEND): """Sort boxes and associated fields according to a scalar field. A common use case is reordering the boxes according to descending scores. @@ -165,9 +154,7 @@ def sort_by_field( order. """ return box_list_to_box_mask_list( - np_box_list_ops.sort_by_field( - boxlist=box_mask_list, field=field, order=order - ) + np_box_list_ops.sort_by_field(boxlist=box_mask_list, field=field, order=order) ) @@ -297,7 +284,7 @@ def multi_class_non_max_suppression( elif len(scores.shape) == 2: if scores.shape[1] is None: raise ValueError( - "scores field must have statically defined second " "dimension" + "scores field must have statically defined second dimension" ) else: raise ValueError("scores field must be of rank 1 or 2") @@ -350,12 +337,8 @@ def prune_non_overlapping_masks(box_mask_list1, box_mask_list2, minoverlap=0.0): Returns: A pruned box_mask_list with size [N', 4]. """ - intersection_over_area = ioa( - box_mask_list2, box_mask_list1 - ) # [M, N] tensor - intersection_over_area = np.amax( - intersection_over_area, axis=0 - ) # [N] tensor + intersection_over_area = ioa(box_mask_list2, box_mask_list1) # [M, N] tensor + intersection_over_area = np.amax(intersection_over_area, axis=0) # [N] tensor keep_bool = np.greater_equal(intersection_over_area, np.array(minoverlap)) keep_inds = np.nonzero(keep_bool)[0] new_box_mask_list1 = gather(box_mask_list1, keep_inds) @@ -419,10 +402,9 @@ def filter_scores_greater_than(box_mask_list, thresh): raise ValueError("Scores should have rank 1 or 2") if len(scores.shape) == 2 and scores.shape[1] != 1: raise ValueError( - "Scores should have rank 1 or have shape " - "consistent with [None, 1]" + "Scores should have rank 1 or have shape consistent with [None, 1]" ) - high_score_indices = np.reshape( - np.where(np.greater(scores, thresh)), [-1] - ).astype(np.int32) + high_score_indices = np.reshape(np.where(np.greater(scores, thresh)), [-1]).astype( + np.int32 + ) return gather(box_mask_list, high_score_indices) diff --git a/slowfast/utils/ava_evaluation/np_box_ops.py b/ava_evaluation/np_box_ops.py similarity index 93% rename from slowfast/utils/ava_evaluation/np_box_ops.py rename to ava_evaluation/np_box_ops.py index 7cadc80f7..d7b337246 100644 --- a/slowfast/utils/ava_evaluation/np_box_ops.py +++ b/ava_evaluation/np_box_ops.py @@ -19,12 +19,9 @@ * Areas: compute bounding box areas * IOU: pairwise intersection-over-union scores """ -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) + +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np @@ -81,11 +78,7 @@ def iou(boxes1, boxes2): intersect = intersection(boxes1, boxes2) area1 = area(boxes1) area2 = area(boxes2) - union = ( - np.expand_dims(area1, axis=1) - + np.expand_dims(area2, axis=0) - - intersect - ) + union = np.expand_dims(area1, axis=1) + np.expand_dims(area2, axis=0) - intersect return intersect / union diff --git a/slowfast/utils/ava_evaluation/np_mask_ops.py b/ava_evaluation/np_mask_ops.py similarity index 92% rename from slowfast/utils/ava_evaluation/np_mask_ops.py rename to ava_evaluation/np_mask_ops.py index 64e36b59b..fb91483b8 100644 --- a/slowfast/utils/ava_evaluation/np_mask_ops.py +++ b/ava_evaluation/np_mask_ops.py @@ -19,12 +19,9 @@ * Areas: compute mask areas * IOU: pairwise intersection-over-union scores """ -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) + +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np EPSILON = 1e-7 @@ -70,9 +67,7 @@ def intersection(masks1, masks2): answer = np.zeros([n, m], dtype=np.float32) for i in np.arange(n): for j in np.arange(m): - answer[i, j] = np.sum( - np.minimum(masks1[i], masks2[j]), dtype=np.float32 - ) + answer[i, j] = np.sum(np.minimum(masks1[i], masks2[j]), dtype=np.float32) return answer @@ -96,11 +91,7 @@ def iou(masks1, masks2): intersect = intersection(masks1, masks2) area1 = area(masks1) area2 = area(masks2) - union = ( - np.expand_dims(area1, axis=1) - + np.expand_dims(area2, axis=0) - - intersect - ) + union = np.expand_dims(area1, axis=1) + np.expand_dims(area2, axis=0) - intersect return intersect / np.maximum(union, EPSILON) diff --git a/slowfast/utils/ava_evaluation/object_detection_evaluation.py b/ava_evaluation/object_detection_evaluation.py similarity index 94% rename from slowfast/utils/ava_evaluation/object_detection_evaluation.py rename to ava_evaluation/object_detection_evaluation.py index 481057200..05e40f9d4 100644 --- a/slowfast/utils/ava_evaluation/object_detection_evaluation.py +++ b/ava_evaluation/object_detection_evaluation.py @@ -27,21 +27,18 @@ Note: This module operates on numpy boxes and box lists. """ -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) +from __future__ import absolute_import, division, print_function, unicode_literals + import collections import logging -import numpy as np from abc import ABCMeta, abstractmethod +import numpy as np + from . import label_map_util, metrics, per_image_evaluation, standard_fields -class DetectionEvaluator(object): +class DetectionEvaluator: """Interface for object detection evalution classes. Example usage of the Evaluator: @@ -182,9 +179,7 @@ def add_single_ground_truth_image_info(self, image_id, groundtruth_dict): raise ValueError("Image with id {} already added.".format(image_id)) groundtruth_classes = ( - groundtruth_dict[ - standard_fields.InputDataFields.groundtruth_classes - ] + groundtruth_dict[standard_fields.InputDataFields.groundtruth_classes] - self._label_id_offset ) # If the key is not present in the groundtruth_dict or the array is empty @@ -216,9 +211,7 @@ def add_single_ground_truth_image_info(self, image_id, groundtruth_dict): standard_fields.InputDataFields.groundtruth_instance_masks not in groundtruth_dict ): - raise ValueError( - "Instance masks not in groundtruth dictionary." - ) + raise ValueError("Instance masks not in groundtruth dictionary.") groundtruth_masks = groundtruth_dict[ standard_fields.InputDataFields.groundtruth_instance_masks ] @@ -255,9 +248,7 @@ def add_single_detected_image_info(self, image_id, detections_dict): ValueError: If detection masks are not in detections dictionary. """ detection_classes = ( - detections_dict[ - standard_fields.DetectionResultFields.detection_classes - ] + detections_dict[standard_fields.DetectionResultFields.detection_classes] - self._label_id_offset ) detection_masks = None @@ -266,9 +257,7 @@ def add_single_detected_image_info(self, image_id, detections_dict): standard_fields.DetectionResultFields.detection_masks not in detections_dict ): - raise ValueError( - "Detection masks not in detections dictionary." - ) + raise ValueError("Detection masks not in detections dictionary.") detection_masks = detections_dict[ standard_fields.DetectionResultFields.detection_masks ] @@ -307,16 +296,12 @@ def evaluate(self): ) = self._evaluation.evaluate() pascal_metrics = { self._metric_prefix - + "Precision/mAP@{}IOU".format( - self._matching_iou_threshold - ): mean_ap + + "Precision/mAP@{}IOU".format(self._matching_iou_threshold): mean_ap } if self._evaluate_corlocs: pascal_metrics[ self._metric_prefix - + "Precision/meanCorLoc@{}IOU".format( - self._matching_iou_threshold - ) + + "Precision/meanCorLoc@{}IOU".format(self._matching_iou_threshold) ] = mean_corloc category_index = label_map_util.create_category_index(self._categories) for idx in range(per_class_ap.size): @@ -437,9 +422,7 @@ class OpenImagesDetectionEvaluator(ObjectDetectionEvaluator): handles those boxes appropriately. """ - def __init__( - self, categories, matching_iou_threshold=0.5, evaluate_corlocs=False - ): + def __init__(self, categories, matching_iou_threshold=0.5, evaluate_corlocs=False): """Constructor. Args: @@ -480,9 +463,7 @@ def add_single_ground_truth_image_info(self, image_id, groundtruth_dict): raise ValueError("Image with id {} already added.".format(image_id)) groundtruth_classes = ( - groundtruth_dict[ - standard_fields.InputDataFields.groundtruth_classes - ] + groundtruth_dict[standard_fields.InputDataFields.groundtruth_classes] - self._label_id_offset ) # If the key is not present in the groundtruth_dict or the array is empty @@ -531,7 +512,7 @@ def add_single_ground_truth_image_info(self, image_id, groundtruth_dict): ) -class ObjectDetectionEvaluation(object): +class ObjectDetectionEvaluation: """Internal implementation of Pascal object detection metrics.""" def __init__( @@ -544,9 +525,7 @@ def __init__( label_id_offset=0, ): if num_groundtruth_classes < 1: - raise ValueError( - "Need at least 1 groundtruth class for evaluation." - ) + raise ValueError("Need at least 1 groundtruth class for evaluation.") self.per_image_eval = per_image_evaluation.PerImageEvaluation( num_groundtruth_classes=num_groundtruth_classes, @@ -621,15 +600,15 @@ def add_single_ground_truth_image_info( if groundtruth_is_difficult_list is None: num_boxes = groundtruth_boxes.shape[0] groundtruth_is_difficult_list = np.zeros(num_boxes, dtype=bool) - self.groundtruth_is_difficult_list[ - image_key - ] = groundtruth_is_difficult_list.astype(dtype=bool) + self.groundtruth_is_difficult_list[image_key] = ( + groundtruth_is_difficult_list.astype(dtype=bool) + ) if groundtruth_is_group_of_list is None: num_boxes = groundtruth_boxes.shape[0] groundtruth_is_group_of_list = np.zeros(num_boxes, dtype=bool) - self.groundtruth_is_group_of_list[ - image_key - ] = groundtruth_is_group_of_list.astype(dtype=bool) + self.groundtruth_is_group_of_list[image_key] = ( + groundtruth_is_group_of_list.astype(dtype=bool) + ) self._update_ground_truth_statistics( groundtruth_class_labels, @@ -664,9 +643,9 @@ def add_single_detected_image_info( ValueError: if the number of boxes, scores and class labels differ in length. """ - if len(detected_boxes) != len(detected_scores) or len( - detected_boxes - ) != len(detected_class_labels): + if len(detected_boxes) != len(detected_scores) or len(detected_boxes) != len( + detected_class_labels + ): raise ValueError( "detected_boxes, detected_scores and " "detected_class_labels should all have same lengths. Got" @@ -692,9 +671,7 @@ def add_single_detected_image_info( groundtruth_is_difficult_list = self.groundtruth_is_difficult_list[ image_key ] - groundtruth_is_group_of_list = self.groundtruth_is_group_of_list[ - image_key - ] + groundtruth_is_group_of_list = self.groundtruth_is_group_of_list[image_key] else: groundtruth_boxes = np.empty(shape=[0, 4], dtype=float) groundtruth_class_labels = np.array([], dtype=int) @@ -748,8 +725,7 @@ def _update_ground_truth_statistics( for class_index in range(self.num_class): num_gt_instances = np.sum( groundtruth_class_labels[ - ~groundtruth_is_difficult_list - & ~groundtruth_is_group_of_list + ~groundtruth_is_difficult_list & ~groundtruth_is_group_of_list ] == class_index ) @@ -790,9 +766,7 @@ def evaluate(self): tp_fp_labels = np.array([], dtype=bool) else: scores = np.concatenate(self.scores_per_class[class_index]) - tp_fp_labels = np.concatenate( - self.tp_fp_labels_per_class[class_index] - ) + tp_fp_labels = np.concatenate(self.tp_fp_labels_per_class[class_index]) if self.use_weighted_mean_ap: all_scores = np.append(all_scores, scores) all_tp_fp_labels = np.append(all_tp_fp_labels, tp_fp_labels) @@ -803,9 +777,7 @@ def evaluate(self): ) self.precisions_per_class.append(precision) self.recalls_per_class.append(recall) - average_precision = metrics.compute_average_precision( - precision, recall - ) + average_precision = metrics.compute_average_precision(precision, recall) self.average_precision_per_class[class_index] = average_precision self.corloc_per_class = metrics.compute_cor_loc( diff --git a/slowfast/utils/ava_evaluation/per_image_evaluation.py b/ava_evaluation/per_image_evaluation.py similarity index 95% rename from slowfast/utils/ava_evaluation/per_image_evaluation.py rename to ava_evaluation/per_image_evaluation.py index aa5684e7a..fa8262d55 100644 --- a/slowfast/utils/ava_evaluation/per_image_evaluation.py +++ b/ava_evaluation/per_image_evaluation.py @@ -20,23 +20,15 @@ Based on the settings, per image evaluation is either performed on boxes or on object masks. """ -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) + +from __future__ import absolute_import, division, print_function, unicode_literals + import numpy as np -from . import ( - np_box_list, - np_box_list_ops, - np_box_mask_list, - np_box_mask_list_ops, -) +from . import np_box_list, np_box_list_ops -class PerImageEvaluation(object): +class PerImageEvaluation: """Evaluate detection result of a single image.""" def __init__(self, num_groundtruth_classes, matching_iou_threshold=0.5): @@ -185,12 +177,12 @@ def _compute_tp_fp( result_scores = [] result_tp_fp_labels = [] for i in range(self.num_groundtruth_classes): - groundtruth_is_difficult_list_at_ith_class = ( - groundtruth_is_difficult_list[groundtruth_class_labels == i] - ) - groundtruth_is_group_of_list_at_ith_class = ( - groundtruth_is_group_of_list[groundtruth_class_labels == i] - ) + groundtruth_is_difficult_list_at_ith_class = groundtruth_is_difficult_list[ + groundtruth_class_labels == i + ] + groundtruth_is_group_of_list_at_ith_class = groundtruth_is_group_of_list[ + groundtruth_class_labels == i + ] ( gt_boxes_at_ith_class, gt_masks_at_ith_class, @@ -329,9 +321,9 @@ def _compute_tp_fp_for_single_class( # Tp-fp evaluation for non-group of boxes (if any). if iou.shape[1] > 0: - groundtruth_nongroup_of_is_difficult_list = ( - groundtruth_is_difficult_list[~groundtruth_is_group_of_list] - ) + groundtruth_nongroup_of_is_difficult_list = groundtruth_is_difficult_list[ + ~groundtruth_is_group_of_list + ] max_overlap_gt_ids = np.argmax(iou, axis=1) is_gt_box_detected = np.zeros(iou.shape[1], dtype=bool) for i in range(num_detected_boxes): @@ -346,9 +338,7 @@ def _compute_tp_fp_for_single_class( return ( scores[~is_matched_to_difficult_box & ~is_matched_to_group_of_box], - tp_fp_labels[ - ~is_matched_to_difficult_box & ~is_matched_to_group_of_box - ], + tp_fp_labels[~is_matched_to_difficult_box & ~is_matched_to_group_of_box], ) def _get_ith_class_arrays( diff --git a/slowfast/utils/ava_evaluation/standard_fields.py b/ava_evaluation/standard_fields.py similarity index 97% rename from slowfast/utils/ava_evaluation/standard_fields.py rename to ava_evaluation/standard_fields.py index 5b1a3df59..0013a1d5e 100644 --- a/slowfast/utils/ava_evaluation/standard_fields.py +++ b/ava_evaluation/standard_fields.py @@ -23,16 +23,10 @@ TfExampleFields: standard fields for tf-example data format (go/tf-example). """ +from __future__ import absolute_import, division, print_function, unicode_literals -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) - -class InputDataFields(object): +class InputDataFields: """Names for the input tensors. Holds the standard data field names to use for identifying input tensors. This @@ -97,7 +91,7 @@ class InputDataFields(object): true_image_shape = "true_image_shape" -class DetectionResultFields(object): +class DetectionResultFields: """Naming conventions for storing the output of the detector. Attributes: @@ -123,7 +117,7 @@ class DetectionResultFields(object): num_detections = "num_detections" -class BoxListFields(object): +class BoxListFields: """Naming conventions for BoxLists. Attributes: @@ -149,7 +143,7 @@ class BoxListFields(object): keypoint_heatmaps = "keypoint_heatmaps" -class TfExampleFields(object): +class TfExampleFields: """TF-example proto feature names for object detection. Holds the standard feature names to load from an Example proto for object diff --git a/configs/ImageNet/MVITv2_B.yaml b/configs/ImageNet/MVITv2_B.yaml new file mode 100644 index 000000000..5d0a330e7 --- /dev/null +++ b/configs/ImageNet/MVITv2_B.yaml @@ -0,0 +1,88 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True +DATA: + # PATH_TO_DATA_DIR: path-to-imagenet-dir + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: False + PATCH_KERNEL: [7, 7] + PATCH_STRIDE: [4, 4] + PATCH_PADDING: [3, 3] + EMBED_DIM: 96 + NUM_HEADS: 1 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.3 + DEPTH: 24 + NORM: "layernorm" + DIM_MUL: [[2, 2.0], [5, 2.0], [21, 2.0]] + HEAD_MUL: [[2, 2.0], [5, 2.0], [21, 2.0]] + POOL_KVQ_KERNEL: [1, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 4, 4] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 1, 1], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 2, 2], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1], [16, 1, 1, 1], [17, 1, 1, 1], [18, 1, 1, 1], [19, 1, 1, 1], [20, 1, 1, 1], [21, 1, 2, 2], [22, 1, 1, 1], [23, 1, 1, 1]] + RESIDUAL_POOLING: True + USE_ABS_POS: False + REL_POS_SPATIAL: True + DIM_MUL_IN_ATT: True +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m9-n6-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.00025 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 70.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 1.0 +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + NUM_ENSEMBLE_VIEWS: 1 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/ImageNet/MVITv2_S.yaml b/configs/ImageNet/MVITv2_S.yaml new file mode 100644 index 000000000..ec5ec7eb2 --- /dev/null +++ b/configs/ImageNet/MVITv2_S.yaml @@ -0,0 +1,88 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True +DATA: + # PATH_TO_DATA_DIR: path-to-imagenet-dir + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: False + PATCH_KERNEL: [7, 7] + PATCH_STRIDE: [4, 4] + PATCH_PADDING: [3, 3] + EMBED_DIM: 96 + NUM_HEADS: 1 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.1 + DEPTH: 16 + NORM: "layernorm" + DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + POOL_KVQ_KERNEL: [1, 3, 3] + POOL_KV_STRIDE: [[0, 1, 4, 4], [1, 1, 2, 2], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1]] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 2, 2], [2, 1, 1, 1], [3, 1, 2, 2], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 2, 2], [15, 1, 1, 1]] + RESIDUAL_POOLING: True + USE_ABS_POS: False + REL_POS_SPATIAL: True + DIM_MUL_IN_ATT: True +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m9-n6-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.00025 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 70.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 1.0 +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + NUM_ENSEMBLE_VIEWS: 1 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/ImageNet/MVITv2_T.yaml b/configs/ImageNet/MVITv2_T.yaml new file mode 100644 index 000000000..6cb7989fe --- /dev/null +++ b/configs/ImageNet/MVITv2_T.yaml @@ -0,0 +1,88 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True +DATA: + # PATH_TO_DATA_DIR: path-to-imagenet-dir + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: False + PATCH_KERNEL: [7, 7] + PATCH_STRIDE: [4, 4] + PATCH_PADDING: [3, 3] + EMBED_DIM: 96 + NUM_HEADS: 1 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.1 + DEPTH: 10 + NORM: "layernorm" + DIM_MUL: [[1, 2.0], [3, 2.0], [8, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [8, 2.0]] + POOL_KVQ_KERNEL: [1, 3, 3] + POOL_KV_STRIDE: [[0, 1, 4, 4], [1, 1, 2, 2], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1]] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 2, 2], [2, 1, 1, 1], [3, 1, 2, 2], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 2, 2], [9, 1, 1, 1]] + RESIDUAL_POOLING: True + USE_ABS_POS: False + REL_POS_SPATIAL: True + DIM_MUL_IN_ATT: True +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m9-n6-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.00025 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 70.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 1.0 +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + NUM_ENSEMBLE_VIEWS: 1 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/ImageNet/REV_MVIT_B_16_CONV.yaml b/configs/ImageNet/REV_MVIT_B_16_CONV.yaml new file mode 100644 index 000000000..0572d92af --- /dev/null +++ b/configs/ImageNet/REV_MVIT_B_16_CONV.yaml @@ -0,0 +1,90 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True + +DATA: + # PATH_TO_DATA_DIR: path-to-imagenet-dir + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: False + PATCH_KERNEL: [7, 7] + PATCH_STRIDE: [4, 4] + PATCH_PADDING: [3, 3] + EMBED_DIM: 96 + NUM_HEADS: 1 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.1 + DEPTH: 16 + NORM: "layernorm" + DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + POOL_KVQ_KERNEL: [1, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 4, 4] + POOL_Q_STRIDE: [[1, 1, 2, 2], [3, 1, 2, 2], [14, 1, 2, 2]] + SEPARATE_QKV : True + REV: + ENABLE: True + RESPATH_FUSE: "concat" + BUFFER_LAYERS : [1,3, 14] + RES_PATH : "conv" + PRE_Q_FUSION: "concat_linear_2" +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m9-n6-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.00025 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 70.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 1.0 +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/ImageNet/REV_VIT_B.yaml b/configs/ImageNet/REV_VIT_B.yaml new file mode 100644 index 000000000..953dff573 --- /dev/null +++ b/configs/ImageNet/REV_VIT_B.yaml @@ -0,0 +1,88 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 512 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True + +DATA: + # PATH_TO_DATA_DIR: path-to-imagenet-dir + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +MVIT: + # USE_SOFTMAX: False + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: False + PATCH_KERNEL: [16, 16] + PATCH_STRIDE: [16, 16] + PATCH_PADDING: [0, 0] + EMBED_DIM: 768 + NUM_HEADS: 12 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.2 # 0.1 + DEPTH: 12 + NORM: "layernorm" + POOL_KV_STRIDE: [] + SEPARATE_QKV : True + REV: + ENABLE: True + RESPATH_FUSE: "concat" + RES_PATH: "conv" + PRE_Q_FUSION: "concat" + +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m6-n5-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0005 + # BASE_LR: 0.0003 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.1 + WARMUP_EPOCHS: 70.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 2.0 +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/ImageNet/REV_VIT_S.yaml b/configs/ImageNet/REV_VIT_S.yaml new file mode 100644 index 000000000..86d867369 --- /dev/null +++ b/configs/ImageNet/REV_VIT_S.yaml @@ -0,0 +1,84 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 512 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True + +DATA: + # PATH_TO_DATA_DIR: path-to-imagenet-dir + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: False + PATCH_KERNEL: [16, 16] + PATCH_STRIDE: [16, 16] + PATCH_PADDING: [0, 0] + EMBED_DIM: 384 + NUM_HEADS: 6 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.1 + DEPTH: 12 + NORM: "layernorm" + POOL_KV_STRIDE: [] + SEPARATE_QKV : True + REV: + ENABLE: True + RESPATH_FUSE: "concat" + RES_PATH: "conv" +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m7-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0005 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 70.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 1.0 # this is commented now +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/Kinetics/MVIT_B_16x4_CONV.yaml b/configs/Kinetics/MVIT_B_16x4_CONV.yaml index b5faead29..fe6a49f55 100644 --- a/configs/Kinetics/MVIT_B_16x4_CONV.yaml +++ b/configs/Kinetics/MVIT_B_16x4_CONV.yaml @@ -79,6 +79,7 @@ MODEL: LOSS_FUNC: soft_cross_entropy DROPOUT_RATE: 0.5 TEST: + ENABLE: True DATASET: kinetics BATCH_SIZE: 64 diff --git a/configs/Kinetics/REV_MVIT_B_16x4_CONV.yaml b/configs/Kinetics/REV_MVIT_B_16x4_CONV.yaml new file mode 100644 index 000000000..af2c72805 --- /dev/null +++ b/configs/Kinetics/REV_MVIT_B_16x4_CONV.yaml @@ -0,0 +1,102 @@ +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 16 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 10 + AUTO_RESUME: True + +DATA: + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + # PATH_TO_DATA_DIR: path-to-k400-dir + TRAIN_JITTER_SCALES_RELATIVE: [0.08, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + +MVIT: + ZERO_DECAY_POS_CLS: False + CLS_EMBED_ON: True + SEP_POS_EMBED: True + USE_ABS_POS: True + DEPTH: 16 + NUM_HEADS: 1 + EMBED_DIM: 96 + PATCH_KERNEL: (3, 7, 7) + PATCH_STRIDE: (2, 4, 4) + PATCH_PADDING: (1, 3, 3) + MLP_RATIO: 4.0 + QKV_BIAS: False + DROPPATH_RATE: 0.05 + NORM: "layernorm" + MODE: "conv" + DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + POOL_KVQ_KERNEL: [3, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + POOL_Q_STRIDE: [[1, 1, 2, 2], [3, 1, 2, 2], [14, 1, 2, 2]] + + REV: + ENABLE: True + RESPATH_FUSE: "concat" + BUFFER_LAYERS : [1, 3, 14] + RES_PATH : "conv" + +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m7-n4-mstd0.5-inc1 + INTERPOLATION: bicubic + NUM_SAMPLE: 2 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False + +MIXUP: + ENABLE: True + ALPHA: 0.8 + LABEL_SMOOTH_VALUE: 0.1 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0001 + LR_POLICY: cosine + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + MAX_EPOCH: 200 + WARMUP_EPOCHS: 50.0 + MOMENTUM: 0.9 + WEIGHT_DECAY: 7e-2 + ZERO_WD_1D_PARAM: True + WARMUP_START_LR: 0.000001 + OPTIMIZING_METHOD: adamw + +MODEL: + NUM_CLASSES: 400 + ARCH: slow + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.5 + +TEST: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 + NUM_ENSEMBLE_VIEWS: 5 # [1, 3, 5, 7, 10] + +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True + +NUM_GPUS: 8 +NUM_SHARDS: 16 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/SSv2/MVITv2_B_32x3.yaml b/configs/SSv2/MVITv2_B_32x3.yaml new file mode 100644 index 000000000..ff0db83de --- /dev/null +++ b/configs/SSv2/MVITv2_B_32x3.yaml @@ -0,0 +1,102 @@ +TRAIN: + ENABLE: True + DATASET: ssv2 + BATCH_SIZE: 16 + EVAL_PERIOD: 2 + CHECKPOINT_PERIOD: 2 + AUTO_RESUME: True + # CHECKPOINT_FILE_PATH: + CHECKPOINT_TYPE: pytorch + CHECKPOINT_EPOCH_RESET: True +DATA: + DECODING_BACKEND: torchvision + NUM_FRAMES: 32 + SAMPLING_RATE: 3 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + INV_UNIFORM_SAMPLE: True + RANDOM_FLIP: False + REVERSE_INPUT_CHANNEL: True + # PATH_TO_DATA_DIR: +MVIT: + ZERO_DECAY_POS_CLS: False + USE_ABS_POS: False + REL_POS_SPATIAL: True + REL_POS_TEMPORAL: True + DEPTH: 24 + NUM_HEADS: 1 + EMBED_DIM: 96 + PATCH_KERNEL: (3, 7, 7) + PATCH_STRIDE: (2, 4, 4) + PATCH_PADDING: (1, 3, 3) + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.4 + NORM: "layernorm" + MODE: "conv" + CLS_EMBED_ON: True + DIM_MUL: [[2, 2.0], [5, 2.0], [21, 2.0]] + HEAD_MUL: [[2, 2.0], [5, 2.0], [21, 2.0]] + POOL_KVQ_KERNEL: [3, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 1, 1], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 2, 2], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1], [16, 1, 1, 1], [17, 1, 1, 1], [18, 1, 1, 1], [19, 1, 1, 1], [20, 1, 1, 1], [21, 1, 2, 2], [22, 1, 1, 1], [23, 1, 1, 1]] + DROPOUT_RATE: 0.0 + DIM_MUL_IN_ATT: True + RESIDUAL_POOLING: True +AUG: + NUM_SAMPLE: 1 + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m7-n4-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +BN: + USE_PRECISE_STATS: False + NUM_BATCHES_PRECISE: 200 +SOLVER: + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0015 + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + WARMUP_START_LR: 1e-4 + WARMUP_EPOCHS: 3.0 + LR_POLICY: cosine + MAX_EPOCH: 100 + MOMENTUM: 0.9 + WEIGHT_DECAY: 1e-4 + OPTIMIZING_METHOD: sgd + COSINE_AFTER_WARMUP: True +MODEL: + NUM_CLASSES: 174 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.5 +TEST: + ENABLE: True + DATASET: ssv2 + BATCH_SIZE: 64 + NUM_ENSEMBLE_VIEWS: 1 + NUM_SPATIAL_CROPS: 3 +DATA_LOADER: + NUM_WORKERS: 2 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +PREFETCH: + NUM_LOADERS: 3 diff --git a/configs/SSv2/MVITv2_L_40x3.yaml b/configs/SSv2/MVITv2_L_40x3.yaml new file mode 100644 index 000000000..a3e918d19 --- /dev/null +++ b/configs/SSv2/MVITv2_L_40x3.yaml @@ -0,0 +1,118 @@ +TRAIN: + ENABLE: True + DATASET: ssv2 + BATCH_SIZE: 8 + EVAL_PERIOD: 2 + CHECKPOINT_PERIOD: 1 + AUTO_RESUME: True + #CHECKPOINT_FILE_PATH: + CHECKPOINT_IN_INIT: True + CHECKPOINT_TYPE: pytorch + CHECKPOINT_EPOCH_RESET: True +DATA: + DECODING_BACKEND: torchvision + NUM_FRAMES: 40 + SAMPLING_RATE: 3 + TRAIN_JITTER_SCALES: [356, 446] + TRAIN_CROP_SIZE: 312 + TEST_CROP_SIZE: 312 + INPUT_CHANNEL_NUM: [3] + INV_UNIFORM_SAMPLE: True + RANDOM_FLIP: False + REVERSE_INPUT_CHANNEL: True + # PATH_TO_DATA_DIR: + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] +MVIT: + ZERO_DECAY_POS_CLS: False + USE_ABS_POS: False + REL_POS_SPATIAL: True + REL_POS_TEMPORAL: True + DEPTH: 48 + NUM_HEADS: 2 + EMBED_DIM: 144 + PATCH_KERNEL: (3, 7, 7) + PATCH_STRIDE: (2, 4, 4) + PATCH_PADDING: (1, 3, 3) + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.75 + NORM: "layernorm" + MODE: "conv" + CLS_EMBED_ON: True + DIM_MUL: [[2, 2.0], [8, 2.0], [44, 2.0]] + HEAD_MUL: [[2, 2.0], [8, 2.0], [44, 2.0]] + POOL_KVQ_KERNEL: [3, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [2, 4, 4] + DROPOUT_RATE: 0.0 + POOL_Q_STRIDE: [[2, 1, 2, 2], [8, 1, 2, 2], [44, 1, 2, 2]] + RESIDUAL_POOLING: True + MODE: "conv" + DIM_MUL_IN_ATT: True + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 1, 1], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 2, 2], [9, 1, 1, 1], [10, 1, 1, 1], + [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1], [16, 1, 1, 1], [17, 1, 1, 1], [18, 1, 1, 1], [19, 1, 1, 1], [20, 1, 1, 1], + [21, 1, 1, 1], [22, 1, 1, 1], [23, 1, 1, 1], [24, 1, 1, 1], [25, 1, 1, 1], [26, 1, 1, 1], [27, 1, 1, 1], [28, 1, 1, 1], [29, 1, 1, 1], [30, 1, 1, 1], + [31, 1, 1, 1], [32, 1, 1, 1], [33, 1, 1, 1], [34, 1, 1, 1], [35, 1, 1, 1], [36, 1, 1, 1], [37, 1, 1, 1], [38, 1, 1, 1], [39, 1, 1, 1], [40, 1, 1, 1], + [41, 1, 1, 1], [42, 1, 1, 1], [43, 1, 1, 1], [44, 1, 2, 2], [45, 1, 1, 1], [46, 1, 1, 1], [47, 1, 1, 1] ] + POOL_FIRST: False # default: False + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + +AUG: + NUM_SAMPLE: 1 + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m7-n4-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 0.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +BN: + USE_PRECISE_STATS: False + NUM_BATCHES_PRECISE: 200 +SOLVER: + CLIP_GRAD_L2NORM: 2.0 + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.00125 + # BASE_LR: 0.0001 + + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + WARMUP_START_LR: 1e-6 + WARMUP_EPOCHS: 3.0 + LR_POLICY: cosine + MAX_EPOCH: 40 + MOMENTUM: 0.9 + WEIGHT_DECAY: 1e-4 + OPTIMIZING_METHOD: sgd + COSINE_AFTER_WARMUP: True +MODEL: + NUM_CLASSES: 174 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.5 + ACT_CHECKPOINT: True +TEST: + ENABLE: True + DATASET: ssv2 + BATCH_SIZE: 8 + NUM_ENSEMBLE_VIEWS: 1 + NUM_SPATIAL_CROPS: 3 +DATA_LOADER: + NUM_WORKERS: 2 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +PREFETCH: + NUM_LOADERS: 3 diff --git a/configs/SSv2/MVITv2_S_16x4.yaml b/configs/SSv2/MVITv2_S_16x4.yaml new file mode 100644 index 000000000..5fd28fe41 --- /dev/null +++ b/configs/SSv2/MVITv2_S_16x4.yaml @@ -0,0 +1,99 @@ +TRAIN: + ENABLE: True + DATASET: ssv2 + BATCH_SIZE: 16 + EVAL_PERIOD: 2 + CHECKPOINT_PERIOD: 2 + AUTO_RESUME: True + # CHECKPOINT_FILE_PATH: + CHECKPOINT_TYPE: pytorch + CHECKPOINT_EPOCH_RESET: True +DATA: + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + INV_UNIFORM_SAMPLE: True + RANDOM_FLIP: False + REVERSE_INPUT_CHANNEL: True + #PATH_TO_DATA_DIR: +MVIT: + ZERO_DECAY_POS_CLS: False + USE_ABS_POS: False + REL_POS_SPATIAL: True + REL_POS_TEMPORAL: True + DEPTH: 16 + NUM_HEADS: 1 + EMBED_DIM: 96 + PATCH_KERNEL: (3, 7, 7) + PATCH_STRIDE: (2, 4, 4) + PATCH_PADDING: (1, 3, 3) + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.4 + NORM: "layernorm" + MODE: "conv" + CLS_EMBED_ON: True + DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + POOL_KVQ_KERNEL: [3, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 2, 2], [2, 1, 1, 1], [3, 1, 2, 2], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 2, 2], [15, 1, 1, 1]] + DROPOUT_RATE: 0.0 + DIM_MUL_IN_ATT: True + RESIDUAL_POOLING: True +AUG: + NUM_SAMPLE: 1 + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m7-n4-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0025 + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + WARMUP_START_LR: 1e-4 + WARMUP_EPOCHS: 3.0 + LR_POLICY: cosine + MAX_EPOCH: 100 + MOMENTUM: 0.9 + WEIGHT_DECAY: 1e-4 + OPTIMIZING_METHOD: sgd + COSINE_AFTER_WARMUP: True +MODEL: + NUM_CLASSES: 174 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.5 +TEST: + ENABLE: True + DATASET: ssv2 + BATCH_SIZE: 64 + NUM_ENSEMBLE_VIEWS: 1 + NUM_SPATIAL_CROPS: 3 +DATA_LOADER: + NUM_WORKERS: 2 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +PREFETCH: + NUM_LOADERS: 3 diff --git a/configs/ssl/BYOL_SlowR50_8x8.yaml b/configs/contrastive_ssl/BYOL_SlowR50_8x8.yaml similarity index 98% rename from configs/ssl/BYOL_SlowR50_8x8.yaml rename to configs/contrastive_ssl/BYOL_SlowR50_8x8.yaml index 0570f0fe2..64efa45df 100644 --- a/configs/ssl/BYOL_SlowR50_8x8.yaml +++ b/configs/contrastive_ssl/BYOL_SlowR50_8x8.yaml @@ -13,7 +13,7 @@ MODEL: DROPOUT_RATE: 0.0 HEAD_ACT: none CONTRASTIVE: - T: 0.5 # default 0.07 + T: 0.5 DIM: 256 # 128 default, if changed, change nCls too NUM_MLP_LAYERS: 2 # default 1 BN_MLP: True diff --git a/configs/ssl/MoCo_SlowR50_8x8.yaml b/configs/contrastive_ssl/MoCo_SlowR50_8x8.yaml similarity index 98% rename from configs/ssl/MoCo_SlowR50_8x8.yaml rename to configs/contrastive_ssl/MoCo_SlowR50_8x8.yaml index 8fb40e07e..09879627d 100644 --- a/configs/ssl/MoCo_SlowR50_8x8.yaml +++ b/configs/contrastive_ssl/MoCo_SlowR50_8x8.yaml @@ -15,7 +15,7 @@ MODEL: DROPOUT_RATE: 0.0 HEAD_ACT: none CONTRASTIVE: - T: 0.1 # default 0.07 + T: 0.1 NUM_MLP_LAYERS: 3 # default 1 SEQUENTIAL: True # def fault MOCO_MULTI_VIEW_QUEUE: True # default: False diff --git a/configs/ssl/SimCLR_SlowR50_8x8.yaml b/configs/contrastive_ssl/SimCLR_SlowR50_8x8.yaml similarity index 100% rename from configs/ssl/SimCLR_SlowR50_8x8.yaml rename to configs/contrastive_ssl/SimCLR_SlowR50_8x8.yaml diff --git a/configs/ssl/SwAV_Slow_R50_8x8.yaml b/configs/contrastive_ssl/SwAV_Slow_R50_8x8.yaml similarity index 100% rename from configs/ssl/SwAV_Slow_R50_8x8.yaml rename to configs/contrastive_ssl/SwAV_Slow_R50_8x8.yaml diff --git a/configs/ssl/finetune_SSv2_Slow_R50_syn0.yaml b/configs/contrastive_ssl/finetune_SSv2_Slow_R50_syn0.yaml similarity index 100% rename from configs/ssl/finetune_SSv2_Slow_R50_syn0.yaml rename to configs/contrastive_ssl/finetune_SSv2_Slow_R50_syn0.yaml diff --git a/configs/ssl/finetune_SSv2_Slow_R50_syn8.yaml b/configs/contrastive_ssl/finetune_SSv2_Slow_R50_syn8.yaml similarity index 100% rename from configs/ssl/finetune_SSv2_Slow_R50_syn8.yaml rename to configs/contrastive_ssl/finetune_SSv2_Slow_R50_syn8.yaml diff --git a/configs/ssl/finetune_ava_Slow_R50_syn0.yaml b/configs/contrastive_ssl/finetune_ava_Slow_R50_syn0.yaml similarity index 100% rename from configs/ssl/finetune_ava_Slow_R50_syn0.yaml rename to configs/contrastive_ssl/finetune_ava_Slow_R50_syn0.yaml diff --git a/configs/ssl/finetune_ava_Slow_R50_syn8.yaml b/configs/contrastive_ssl/finetune_ava_Slow_R50_syn8.yaml similarity index 100% rename from configs/ssl/finetune_ava_Slow_R50_syn8.yaml rename to configs/contrastive_ssl/finetune_ava_Slow_R50_syn8.yaml diff --git a/configs/ssl/finetune_ucf_Slow_R50_syn0.yaml b/configs/contrastive_ssl/finetune_ucf_Slow_R50_syn0.yaml similarity index 100% rename from configs/ssl/finetune_ucf_Slow_R50_syn0.yaml rename to configs/contrastive_ssl/finetune_ucf_Slow_R50_syn0.yaml diff --git a/configs/ssl/finetune_ucf_Slow_R50_syn8.yaml b/configs/contrastive_ssl/finetune_ucf_Slow_R50_syn8.yaml similarity index 100% rename from configs/ssl/finetune_ucf_Slow_R50_syn8.yaml rename to configs/contrastive_ssl/finetune_ucf_Slow_R50_syn8.yaml diff --git a/configs/ssl/linear_k400_Slow_8x8_R50_syn0.yaml b/configs/contrastive_ssl/linear_k400_Slow_8x8_R50_syn0.yaml similarity index 100% rename from configs/ssl/linear_k400_Slow_8x8_R50_syn0.yaml rename to configs/contrastive_ssl/linear_k400_Slow_8x8_R50_syn0.yaml diff --git a/configs/ssl/linear_k400_Slow_8x8_R50_syn8.yaml b/configs/contrastive_ssl/linear_k400_Slow_8x8_R50_syn8.yaml similarity index 100% rename from configs/ssl/linear_k400_Slow_8x8_R50_syn8.yaml rename to configs/contrastive_ssl/linear_k400_Slow_8x8_R50_syn8.yaml diff --git a/configs/masked_ssl/in1k_VIT_B_MaskFeat_FT.yaml b/configs/masked_ssl/in1k_VIT_B_MaskFeat_FT.yaml new file mode 100644 index 000000000..c8f014f75 --- /dev/null +++ b/configs/masked_ssl/in1k_VIT_B_MaskFeat_FT.yaml @@ -0,0 +1,84 @@ +TASK: ssl_eval +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 512 + EVAL_PERIOD: 10 + CHECKPOINT_PERIOD: 10 + CHECKPOINT_EPOCH_RESET: True + AUTO_RESUME: True +DATA: + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +BN: + USE_PRECISE_STATS: True + NUM_BATCHES_PRECISE: 200 +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: True + PATCH_KERNEL: [16, 16] + PATCH_STRIDE: [16, 16] + PATCH_PADDING: [0, 0] + EMBED_DIM: 768 + NUM_HEADS: 12 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.1 + LAYER_SCALE_INIT_VALUE: 0.0 + DEPTH: 12 + NORM: "layernorm" + HEAD_INIT_SCALE: 0.001 + USE_MEAN_POOLING: True +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m9-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + LAYER_DECAY: 0.65 + BASE_LR: 0.002 + LR_POLICY: cosine + MAX_EPOCH: 100 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 20.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/in1k_VIT_B_MaskFeat_PT.yaml b/configs/masked_ssl/in1k_VIT_B_MaskFeat_PT.yaml new file mode 100644 index 000000000..88e0e8997 --- /dev/null +++ b/configs/masked_ssl/in1k_VIT_B_MaskFeat_PT.yaml @@ -0,0 +1,80 @@ +TASK: ssl +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 10000 + CHECKPOINT_PERIOD: 10 + AUTO_RESUME: True +DATA: + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +BN: + USE_PRECISE_STATS: True + NUM_BATCHES_PRECISE: 200 +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: True + PATCH_KERNEL: [16, 16] + PATCH_STRIDE: [16, 16] + PATCH_PADDING: [0, 0] + EMBED_DIM: 768 + NUM_HEADS: 12 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.0 + DEPTH: 12 + NORM: "layernorm" +MASK: + ENABLE: True + PRETRAIN_DEPTH: [11] + HEAD_TYPE: "separate" + PRED_HOG: True +AUG: + ENABLE: True + RE_PROB: 0.0 + COLOR_JITTER: None + AA_TYPE: "" + GEN_MASK_LOADER: True + MASK_RATIO: 0.4 +MIXUP: + ENABLE: False +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0002 # bs256 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 30.0 + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 0.02 +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/configs/masked_ssl/in1k_VIT_L_MaskFeat_FT.yaml b/configs/masked_ssl/in1k_VIT_L_MaskFeat_FT.yaml new file mode 100644 index 000000000..aa356727f --- /dev/null +++ b/configs/masked_ssl/in1k_VIT_L_MaskFeat_FT.yaml @@ -0,0 +1,84 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 1 + CHECKPOINT_PERIOD: 5 + CHECKPOINT_EPOCH_RESET: True + AUTO_RESUME: True + MIXED_PRECISION: False +DATA: + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +BN: + USE_PRECISE_STATS: True + NUM_BATCHES_PRECISE: 200 +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: True + PATCH_KERNEL: [16, 16] + PATCH_STRIDE: [16, 16] + PATCH_PADDING: [0, 0] + EMBED_DIM: 1024 + NUM_HEADS: 16 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.1 + LAYER_SCALE_INIT_VALUE: 0.0 + DEPTH: 24 + NORM: "layernorm" + HEAD_INIT_SCALE: 0.001 + USE_MEAN_POOLING: True +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m9-mstd0.5-inc1 # ! + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + LAYER_DECAY: 0.75 + BASE_LR: 0.001 + LR_POLICY: cosine + MAX_EPOCH: 50 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 5.0 + WARMUP_START_LR: 1e-8 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True +MODEL: + NUM_CLASSES: 1000 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/in1k_VIT_L_MaskFeat_PT.yaml b/configs/masked_ssl/in1k_VIT_L_MaskFeat_PT.yaml new file mode 100644 index 000000000..e3487ebf1 --- /dev/null +++ b/configs/masked_ssl/in1k_VIT_L_MaskFeat_PT.yaml @@ -0,0 +1,80 @@ +TRAIN: + ENABLE: True + DATASET: imagenet + BATCH_SIZE: 256 + EVAL_PERIOD: 10000 + CHECKPOINT_PERIOD: 10 + AUTO_RESUME: True + MIXED_PRECISION: False +DATA: + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + MEAN: [0.485, 0.456, 0.406] + STD: [0.229, 0.224, 0.225] + NUM_FRAMES: 1 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] +BN: + USE_PRECISE_STATS: True + NUM_BATCHES_PRECISE: 200 +MVIT: + PATCH_2D: True + ZERO_DECAY_POS_CLS: False + MODE: "conv" + CLS_EMBED_ON: True + PATCH_KERNEL: [16, 16] + PATCH_STRIDE: [16, 16] + PATCH_PADDING: [0, 0] + EMBED_DIM: 1024 + NUM_HEADS: 16 + MLP_RATIO: 4.0 + QKV_BIAS: True + DROPPATH_RATE: 0.0 + DEPTH: 24 + NORM: "layernorm" +MASK: + ENABLE: True + PRETRAIN_DEPTH: [23] + HEAD_TYPE: "separate" + PRED_HOG: True +AUG: + ENABLE: True + RE_PROB: 0.0 + COLOR_JITTER: None # important, not using color aug + AA_TYPE: "" # important, not using color aug + GEN_MASK_LOADER: True + MASK_RATIO: 0.4 +MIXUP: + ENABLE: False +SOLVER: + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0002 + LR_POLICY: cosine + MAX_EPOCH: 300 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + WARMUP_EPOCHS: 30.0 + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 0.02 +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: imagenet + BATCH_SIZE: 256 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/configs/masked_ssl/k400_MVITv2_L_16x4_FT.yaml b/configs/masked_ssl/k400_MVITv2_L_16x4_FT.yaml new file mode 100644 index 000000000..38f9118f2 --- /dev/null +++ b/configs/masked_ssl/k400_MVITv2_L_16x4_FT.yaml @@ -0,0 +1,108 @@ +TASK: ssl_eval +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 16 + EVAL_PERIOD: 5 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + CHECKPOINT_EPOCH_RESET: True +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + DECODING_SHORT_SIZE: 320 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + TRAIN_JITTER_SCALES_RELATIVE: [0.08, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True + PATCH_KERNEL: (3, 7, 7) + PATCH_STRIDE: (2, 4, 4) + PATCH_PADDING: (1, 3, 3) + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + EMBED_DIM: 144 + NUM_HEADS: 2 + DEPTH: 48 # [2, 6, 36, 2] + DIM_MUL: [[2, 2.0], [8, 2.0], [44, 2.0]] + HEAD_MUL: [[2, 2.0], [8, 2.0], [44, 2.0]] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 1, 1], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 2, 2], [9, 1, 1, 1], [10, 1, 1, 1], + [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1], [16, 1, 1, 1], [17, 1, 1, 1], [18, 1, 1, 1], [19, 1, 1, 1], [20, 1, 1, 1], + [21, 1, 1, 1], [22, 1, 1, 1], [23, 1, 1, 1], [24, 1, 1, 1], [25, 1, 1, 1], [26, 1, 1, 1], [27, 1, 1, 1], [28, 1, 1, 1], [29, 1, 1, 1], [30, 1, 1, 1], + [31, 1, 1, 1], [32, 1, 1, 1], [33, 1, 1, 1], [34, 1, 1, 1], [35, 1, 1, 1], [36, 1, 1, 1], [37, 1, 1, 1], [38, 1, 1, 1], [39, 1, 1, 1], [40, 1, 1, 1], + [41, 1, 1, 1], [42, 1, 1, 1], [43, 1, 1, 1], [44, 1, 2, 2], [45, 1, 1, 1], [46, 1, 1, 1], [47, 1, 1, 1]] + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + POOL_KVQ_KERNEL: [3, 3, 3] + # Please check if your loaded model has cls token. + # It would be better to have cls token. But some of old models forget to turn this on. + CLS_EMBED_ON: True # default: True + # uncomment comment for abs pos: + # SEP_POS_EMBED: True + + # uncomment for rel pos: + USE_ABS_POS: False # default: True + REL_POS_SPATIAL: True # default: false + REL_POS_TEMPORAL: True # default: false + + MODE: "conv" + RESIDUAL_POOLING: True + DROPPATH_RATE: 0.2 + LAYER_SCALE_INIT_VALUE: 0.0 + USE_MEAN_POOLING: True + HEAD_INIT_SCALE: 0.001 +AUG: + ENABLE: True + COLOR_JITTER: None + AA_TYPE: rand-m7-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + MAX_EPOCH: 75 + LAYER_DECAY: 0.875 + WARMUP_EPOCHS: 5.0 + CLIP_GRAD_L2NORM: 5.0 + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0006 + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-5 + WARMUP_START_LR: 1e-8 + LR_POLICY: cosine + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + OPTIMIZING_METHOD: adamw +MODEL: + NUM_CLASSES: 400 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy # default cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/k400_MVITv2_L_16x4_MaskFeat_PT.yaml b/configs/masked_ssl/k400_MVITv2_L_16x4_MaskFeat_PT.yaml new file mode 100644 index 000000000..0ce804617 --- /dev/null +++ b/configs/masked_ssl/k400_MVITv2_L_16x4_MaskFeat_PT.yaml @@ -0,0 +1,105 @@ +TASK: ssl +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 32 + EVAL_PERIOD: 100000 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + DECODING_SHORT_SIZE: 320 + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True + PATCH_KERNEL: (3, 7, 7) + PATCH_STRIDE: (2, 4, 4) + PATCH_PADDING: (1, 3, 3) + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + EMBED_DIM: 144 + NUM_HEADS: 2 + DEPTH: 48 + DIM_MUL: [[2, 2.0], [8, 2.0], [44, 2.0]] + HEAD_MUL: [[2, 2.0], [8, 2.0], [44, 2.0]] + # Highlight: [44, 1, 1, 1] instead of [44, 1, 2, 2] for 14x14 output + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 1, 1], [2, 1, 2, 2], [3, 1, 1, 1], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 2, 2], [9, 1, 1, 1], [10, 1, 1, 1], + [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1], [16, 1, 1, 1], [17, 1, 1, 1], [18, 1, 1, 1], [19, 1, 1, 1], [20, 1, 1, 1], + [21, 1, 1, 1], [22, 1, 1, 1], [23, 1, 1, 1], [24, 1, 1, 1], [25, 1, 1, 1], [26, 1, 1, 1], [27, 1, 1, 1], [28, 1, 1, 1], [29, 1, 1, 1], [30, 1, 1, 1], + [31, 1, 1, 1], [32, 1, 1, 1], [33, 1, 1, 1], [34, 1, 1, 1], [35, 1, 1, 1], [36, 1, 1, 1], [37, 1, 1, 1], [38, 1, 1, 1], [39, 1, 1, 1], [40, 1, 1, 1], + [41, 1, 1, 1], [42, 1, 1, 1], [43, 1, 1, 1], [44, 1, 1, 1], [45, 1, 1, 1], [46, 1, 1, 1], [47, 1, 1, 1]] + DROPPATH_RATE: 0.0 + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + POOL_KVQ_KERNEL: [3, 3, 3] + CLS_EMBED_ON: True # defauult: True + # uncomment comment for abs pos: + # SEP_POS_EMBED: True + + # uncomment for rel pos: + USE_ABS_POS: False # default: True + REL_POS_SPATIAL: True # default: false + REL_POS_TEMPORAL: True # default: false + + MODE: "conv" + RESIDUAL_POOLING: True +MASK: + ENABLE: True + PRETRAIN_DEPTH: [47] + PRED_HOG: True # default: false +AUG: + ENABLE: True + COLOR_JITTER: None + AA_TYPE: "" + RE_PROB: 0.0 + GEN_MASK_LOADER: True + MASK_RATIO: 0.4 + + # Mask Cube (Default) + MASK_TUBE: False + MASK_FRAMES: False + MASK_WINDOW_SIZE: [8, 7, 7] + +MIXUP: + ENABLE: False +SOLVER: + CLIP_GRAD_L2NORM: 0.02 + BASE_LR: 0.0001 + BASE_LR_SCALE_NUM_SHARDS: True + LR_POLICY: cosine + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + MAX_EPOCH: 300 + WARMUP_EPOCHS: 10.0 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + ZERO_WD_1D_PARAM: True + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/configs/masked_ssl/k400_MVITv2_S_16x4_FT.yaml b/configs/masked_ssl/k400_MVITv2_S_16x4_FT.yaml new file mode 100644 index 000000000..45460b9c7 --- /dev/null +++ b/configs/masked_ssl/k400_MVITv2_S_16x4_FT.yaml @@ -0,0 +1,101 @@ +TASK: ssl_eval +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 16 + EVAL_PERIOD: 10 + CHECKPOINT_EPOCH_RESET: True + CHECKPOINT_PERIOD: 10 + AUTO_RESUME: True + MIXED_PRECISION: False +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + TRAIN_JITTER_SCALES_RELATIVE: [0.08, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + DECODING_SHORT_SIZE: 320 +MVIT: + DEPTH: 16 # def 12 + NUM_HEADS: 3 # default 12 16 + EMBED_DIM: 384 # def 192 768 384 1024 + NUM_HEADS: 1 # default 12 16 + EMBED_DIM: 96 # def 256 192 768 384 1024 + PATCH_KERNEL: [3, 7, 7] + PATCH_STRIDE: [2, 4, 4] # default:empty + PATCH_PADDING: [1, 3, 3] + ZERO_DECAY_POS_CLS: False + QKV_BIAS: True + DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 2, 2], [2, 1, 1, 1], [3, 1, 2, 2], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 2, 2], [15, 1, 1, 1]] + POOL_KVQ_KERNEL: [3, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + CLS_EMBED_ON: True # default: True + USE_ABS_POS: False # default: True + SEP_POS_EMBED: True # default: false + REL_POS_SPATIAL: True # default: false + REL_POS_TEMPORAL: True # default: false + RESIDUAL_POOLING: True + MODE: "conv" + DROPPATH_RATE: 0.1 + LAYER_SCALE_INIT_VALUE: 0.0 + USE_MEAN_POOLING: True + HEAD_INIT_SCALE: 0.001 +AUG: + ENABLE: True + COLOR_JITTER: 0.4 + AA_TYPE: rand-m7-mstd0.5-inc1 + INTERPOLATION: bicubic + NUM_SAMPLE: 2 + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + ZERO_WD_1D_PARAM: True + CLIP_GRAD_L2NORM: 5.0 + LAYER_DECAY: 0.75 # ! + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 0.0006 # ! + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + WARMUP_START_LR: 1e-6 + WARMUP_EPOCHS: 20.0 + LR_POLICY: cosine + MAX_EPOCH: 100 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + OPTIMIZING_METHOD: adamw + COSINE_AFTER_WARMUP: True +MODEL: + NUM_CLASSES: 400 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy + DROPOUT_RATE: 0.0 +TEST: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 + NUM_TEMPORAL_CLIPS: [1, 3, 4, 5, 7, 10] +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/k400_MVITv2_S_16x4_MaskFeat_PT.yaml b/configs/masked_ssl/k400_MVITv2_S_16x4_MaskFeat_PT.yaml new file mode 100644 index 000000000..99388faf8 --- /dev/null +++ b/configs/masked_ssl/k400_MVITv2_S_16x4_MaskFeat_PT.yaml @@ -0,0 +1,98 @@ +TASK: ssl +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 32 + EVAL_PERIOD: 1000000 + CHECKPOINT_PERIOD: 10 + AUTO_RESUME: True +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + DECODING_SHORT_SIZE: 320 + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] +MVIT: + DEPTH: 16 + NUM_HEADS: 1 + EMBED_DIM: 96 + PATCH_KERNEL: [3, 7, 7] + PATCH_STRIDE: [2, 4, 4] + PATCH_PADDING: [1, 3, 3] + ZERO_DECAY_POS_CLS: False + QKV_BIAS: True + DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]] + # Highlight: [14, 1, 1, 1] instead of [14, 1, 2, 2] for 14x14 output + POOL_Q_STRIDE: [[0, 1, 1, 1], [1, 1, 2, 2], [2, 1, 1, 1], [3, 1, 2, 2], [4, 1, 1, 1], [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, 1, 1], [8, 1, 1, 1], [9, 1, 1, 1], [10, 1, 1, 1], [11, 1, 1, 1], [12, 1, 1, 1], [13, 1, 1, 1], [14, 1, 1, 1], [15, 1, 1, 1]] + POOL_KVQ_KERNEL: [3, 3, 3] + POOL_KV_STRIDE_ADAPTIVE: [1, 8, 8] + DROPPATH_RATE: 0.0 + CLS_EMBED_ON: True # default: True + USE_ABS_POS: False # default: True + SEP_POS_EMBED: True # default: false + REL_POS_SPATIAL: True # default: false + REL_POS_TEMPORAL: True # default: false + RESIDUAL_POOLING: True + MODE: "conv" +MASK: + ENABLE: True + PRETRAIN_DEPTH: [15] + HEAD_TYPE: "separate" + PRED_HOG: True +AUG: + ENABLE: True + COLOR_JITTER: None + AA_TYPE: "" + RE_PROB: 0.0 + GEN_MASK_LOADER: True + MASK_RATIO: 0.4 + + # Mask Cube (Default) + MASK_TUBE: False + MASK_FRAMES: False + MASK_WINDOW_SIZE: [8, 7, 7] +MIXUP: + ENABLE: False +BN: + USE_PRECISE_STATS: False + NUM_BATCHES_PRECISE: 200 +SOLVER: + CLIP_GRAD_L2NORM: 0.02 + BASE_LR: 0.0001 + BASE_LR_SCALE_NUM_SHARDS: True + LR_POLICY: cosine + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + MAX_EPOCH: 300 + WARMUP_EPOCHS: 10.0 + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + ZERO_WD_1D_PARAM: True + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: kinetics + BATCH_SIZE: 16 + NUM_SPATIAL_CROPS: 1 + NUM_TEMPORAL_CLIPS: [5, 10] +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/configs/masked_ssl/k400_VIT_B_16x4_FT.yaml b/configs/masked_ssl/k400_VIT_B_16x4_FT.yaml new file mode 100644 index 000000000..ccce0f7e7 --- /dev/null +++ b/configs/masked_ssl/k400_VIT_B_16x4_FT.yaml @@ -0,0 +1,90 @@ +TASK: ssl_eval +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 1 + EVAL_PERIOD: 1 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + CHECKPOINT_EPOCH_RESET: True +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + DECODING_SHORT_SIZE: 320 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + TRAIN_JITTER_SCALES_RELATIVE: [0.08, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True + PATCH_KERNEL: (2, 16, 16) + PATCH_STRIDE: (2, 16, 16) + PATCH_PADDING: (0, 0, 0) + EMBED_DIM: 768 # vit-B, change depth & layerdecay + NUM_HEADS: 12 + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + DEPTH: 12 + MODE: "conv" + DROPPATH_RATE: 0.1 + LAYER_SCALE_INIT_VALUE: 0.0 + USE_MEAN_POOLING: True + HEAD_INIT_SCALE: 0.001 +AUG: + NUM_SAMPLE: 2 + ENABLE: True + COLOR_JITTER: None + AA_TYPE: rand-m7-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + MAX_EPOCH: 75 + LAYER_DECAY: 0.75 + LAYER_DECAY: 0.65 # vit-B + WARMUP_EPOCHS: 5.0 + CLIP_GRAD_L2NORM: 5.0 + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 6e-4 + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-5 + WARMUP_START_LR: 1e-8 + LR_POLICY: cosine + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + OPTIMIZING_METHOD: adamw +MODEL: + NUM_CLASSES: 400 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy # default cross_entropy + DROPOUT_RATE: 0.3 +TEST: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 3 + NUM_TEMPORAL_CLIPS: [5, 7, 10] +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/k400_VIT_B_16x4_MAE_PT.yaml b/configs/masked_ssl/k400_VIT_B_16x4_MAE_PT.yaml new file mode 100644 index 000000000..5114d754d --- /dev/null +++ b/configs/masked_ssl/k400_VIT_B_16x4_MAE_PT.yaml @@ -0,0 +1,86 @@ +TASK: ssl +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + EVAL_PERIOD: 100000 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + KILL_LOSS_EXPLOSION_FACTOR: 2.0 # default 0.0 (not enforced) +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + DECODING_SHORT_SIZE: 320 + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + TRAIN_CROP_NUM_TEMPORAL: 4 # default 1 +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True # default: False needs to be false for fixed sincos + PATCH_KERNEL: (2, 16, 16) + PATCH_STRIDE: (2, 16, 16) + PATCH_PADDING: (0, 0, 0) + EMBED_DIM: 768 # vit-B + NUM_HEADS: 12 + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + DEPTH: 12 + DROPPATH_RATE: 0.0 + MODE: "conv" + CLS_EMBED_ON: True +MASK: + ENABLE: True + MAE_ON: True # default False + MAE_RND_MASK: True # default: false + PRETRAIN_DEPTH: [11] + HEAD_TYPE: "separate_xformer" + DECODER_DEPTH: 4 + DECODER_EMBED_DIM: 512 +AUG: + ENABLE: True + COLOR_JITTER: None + AA_TYPE: "" + RE_PROB: 0.0 + MASK_RATIO: 0.9 +MIXUP: + ENABLE: False +SOLVER: + CLIP_GRAD_L2NORM: 0.02 + BASE_LR: 1e-4 + BASE_LR_SCALE_NUM_SHARDS: True + LR_POLICY: cosine + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + MAX_EPOCH: 200 + WARMUP_EPOCHS: 60.0 + BETAS: (0.9, 0.95) # default (0.9, 0.999) + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + ZERO_WD_1D_PARAM: True # default False + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/configs/masked_ssl/k400_VIT_H_16x4_FT.yaml b/configs/masked_ssl/k400_VIT_H_16x4_FT.yaml new file mode 100644 index 000000000..bacda1650 --- /dev/null +++ b/configs/masked_ssl/k400_VIT_H_16x4_FT.yaml @@ -0,0 +1,88 @@ +TASK: ssl_eval +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 16 + EVAL_PERIOD: 5 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + CHECKPOINT_EPOCH_RESET: True +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + DECODING_SHORT_SIZE: 320 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + TRAIN_JITTER_SCALES_RELATIVE: [0.08, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True + PATCH_KERNEL: (2, 16, 16) + PATCH_STRIDE: (2, 16, 16) + PATCH_PADDING: (0, 0, 0) + EMBED_DIM: 1280 + NUM_HEADS: 16 + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + DEPTH: 32 + MODE: "conv" + DROPPATH_RATE: 0.2 + LAYER_SCALE_INIT_VALUE: 0.0 + USE_MEAN_POOLING: True + HEAD_INIT_SCALE: 0.001 +AUG: + ENABLE: True + COLOR_JITTER: None + AA_TYPE: rand-m7-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + MAX_EPOCH: 75 + LAYER_DECAY: 0.8 + WARMUP_EPOCHS: 5.0 + CLIP_GRAD_L2NORM: 5.0 + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 3e-4 # works best + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + WARMUP_START_LR: 1e-8 + LR_POLICY: cosine + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + OPTIMIZING_METHOD: adamw +MODEL: + NUM_CLASSES: 400 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy # default cross_entropy + DROPOUT_RATE: 0.5 +TEST: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 3 + NUM_TEMPORAL_CLIPS: [5, 7, 10] +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/k400_VIT_H_16x4_MAE_PT.yaml b/configs/masked_ssl/k400_VIT_H_16x4_MAE_PT.yaml new file mode 100644 index 000000000..5e4504612 --- /dev/null +++ b/configs/masked_ssl/k400_VIT_H_16x4_MAE_PT.yaml @@ -0,0 +1,88 @@ +TASK: ssl +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 32 + EVAL_PERIOD: 100000 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + KILL_LOSS_EXPLOSION_FACTOR: 2.0 # default 0.0 (not enforced) +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + DECODING_SHORT_SIZE: 320 + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + TRAIN_CROP_NUM_TEMPORAL: 4 # default 1 +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True + PATCH_KERNEL: (2, 16, 16) + PATCH_STRIDE: (2, 16, 16) + PATCH_PADDING: (0, 0, 0) + EMBED_DIM: 1280 + NUM_HEADS: 16 + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + DEPTH: 32 + DROPPATH_RATE: 0.0 + MODE: "conv" + CLS_EMBED_ON: True # default: True +MASK: + ENABLE: True + MAE_ON: True # default False + MAE_RND_MASK: True # default: false + PRETRAIN_DEPTH: [31] + HEAD_TYPE: "separate_xformer" + DECODER_DEPTH: 4 # default: 0 + DECODER_EMBED_DIM: 512 +AUG: + NUM_SAMPLE: 1 + ENABLE: True + COLOR_JITTER: None + AA_TYPE: "" + RE_PROB: 0.0 + MASK_RATIO: 0.9 + +MIXUP: + ENABLE: False +SOLVER: + CLIP_GRAD_L2NORM: 0.02 + BASE_LR: 1e-4 # for bs32 + BASE_LR_SCALE_NUM_SHARDS: True + LR_POLICY: cosine + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + MAX_EPOCH: 400 + WARMUP_EPOCHS: 60.0 + BETAS: (0.9, 0.95) # default (0.9, 0.999) + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + ZERO_WD_1D_PARAM: True # default False + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/configs/masked_ssl/k400_VIT_L_16x4_FT.yaml b/configs/masked_ssl/k400_VIT_L_16x4_FT.yaml new file mode 100644 index 000000000..68fa2d0ed --- /dev/null +++ b/configs/masked_ssl/k400_VIT_L_16x4_FT.yaml @@ -0,0 +1,89 @@ +TASK: ssl_eval +TRAIN: + DATASET: kinetics + BATCH_SIZE: 16 + EVAL_PERIOD: 5 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + CHECKPOINT_EPOCH_RESET: True +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + DECODING_SHORT_SIZE: 320 + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + TRAIN_JITTER_SCALES_RELATIVE: [0.08, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True + PATCH_KERNEL: (2, 16, 16) + PATCH_STRIDE: (2, 16, 16) + PATCH_PADDING: (0, 0, 0) + EMBED_DIM: 1024 + NUM_HEADS: 16 + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + DEPTH: 24 + MODE: "conv" + DROPPATH_RATE: 0.2 + LAYER_SCALE_INIT_VALUE: 0.0 + USE_MEAN_POOLING: True + HEAD_INIT_SCALE: 0.001 +AUG: + NUM_SAMPLE: 2 + ENABLE: True + COLOR_JITTER: None + AA_TYPE: rand-m7-mstd0.5-inc1 + INTERPOLATION: bicubic + RE_PROB: 0.25 + RE_MODE: pixel + RE_COUNT: 1 + RE_SPLIT: False +MIXUP: + ENABLE: True + ALPHA: 0.8 + CUTMIX_ALPHA: 1.0 + PROB: 1.0 + SWITCH_PROB: 0.5 + LABEL_SMOOTH_VALUE: 0.1 +SOLVER: + MAX_EPOCH: 75 + MAX_EPOCH: 50 + LAYER_DECAY: 0.75 + WARMUP_EPOCHS: 5.0 + CLIP_GRAD_L2NORM: 5.0 + ZERO_WD_1D_PARAM: True + BASE_LR_SCALE_NUM_SHARDS: True + BASE_LR: 3e-4 # works best + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-5 + WARMUP_START_LR: 1e-8 + LR_POLICY: cosine + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + OPTIMIZING_METHOD: adamw +MODEL: + NUM_CLASSES: 400 + ARCH: mvit + MODEL_NAME: MViT + LOSS_FUNC: soft_cross_entropy # default cross_entropy + DROPOUT_RATE: 0.3 +TEST: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 3 + NUM_TEMPORAL_CLIPS: [5, 7, 10] +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . diff --git a/configs/masked_ssl/k400_VIT_L_16x4_MAE_PT.yaml b/configs/masked_ssl/k400_VIT_L_16x4_MAE_PT.yaml new file mode 100644 index 000000000..51bf33f07 --- /dev/null +++ b/configs/masked_ssl/k400_VIT_L_16x4_MAE_PT.yaml @@ -0,0 +1,87 @@ +TASK: ssl +TRAIN: + ENABLE: True + DATASET: kinetics + BATCH_SIZE: 32 # 32 fits 32g + EVAL_PERIOD: 100000 + CHECKPOINT_PERIOD: 5 + AUTO_RESUME: True + KILL_LOSS_EXPLOSION_FACTOR: 2.0 # default 0.0 (not enforced) +DATA: + USE_OFFSET_SAMPLING: True + DECODING_BACKEND: torchvision + NUM_FRAMES: 16 + SAMPLING_RATE: 4 + TRAIN_JITTER_SCALES: [256, 320] + TRAIN_CROP_SIZE: 224 + TEST_CROP_SIZE: 224 + INPUT_CHANNEL_NUM: [3] + DECODING_SHORT_SIZE: 320 + TRAIN_JITTER_SCALES_RELATIVE: [0.5, 1.0] + TRAIN_JITTER_ASPECT_RELATIVE: [0.75, 1.3333] + TRAIN_CROP_NUM_TEMPORAL: 4 # default 1 +MVIT: + ZERO_DECAY_POS_CLS: False + SEP_POS_EMBED: True # default: False needs to be false for fixed sincos + PATCH_KERNEL: (2, 16, 16) + PATCH_STRIDE: (2, 16, 16) + PATCH_PADDING: (0, 0, 0) + EMBED_DIM: 1024 + NUM_HEADS: 16 + MLP_RATIO: 4.0 + QKV_BIAS: True + NORM: "layernorm" + DEPTH: 24 + DROPPATH_RATE: 0.0 + MODE: "conv" + CLS_EMBED_ON: True # default: True +MASK: + ENABLE: True + MAE_ON: True # default False + MAE_RND_MASK: True # default: false + PRETRAIN_DEPTH: [23] + HEAD_TYPE: "separate_xformer" + DECODER_DEPTH: 4 # default: 0 + DECODER_EMBED_DIM: 512 +AUG: + # NUM_SAMPLE: 4 + ENABLE: True + COLOR_JITTER: None + AA_TYPE: "" + RE_PROB: 0.0 + MASK_RATIO: 0.9 +MIXUP: + ENABLE: False +SOLVER: + CLIP_GRAD_L2NORM: 0.02 + BASE_LR: 1e-4 # for bs64 + BASE_LR_SCALE_NUM_SHARDS: True + LR_POLICY: cosine + COSINE_AFTER_WARMUP: True + COSINE_END_LR: 1e-6 + MAX_EPOCH: 200 + WARMUP_EPOCHS: 30.0 + BETAS: (0.9, 0.95) # default (0.9, 0.999) + MOMENTUM: 0.9 + WEIGHT_DECAY: 0.05 + ZERO_WD_1D_PARAM: True # default False + WARMUP_START_LR: 1e-6 + OPTIMIZING_METHOD: adamw +MODEL: + ARCH: maskmvit + MODEL_NAME: MaskMViT + LOSS_FUNC: multi_mse + DROPOUT_RATE: 0.0 +TEST: + ENABLE: False + DATASET: kinetics + BATCH_SIZE: 64 + NUM_SPATIAL_CROPS: 1 +DATA_LOADER: + NUM_WORKERS: 8 + PIN_MEMORY: True +NUM_GPUS: 8 +NUM_SHARDS: 1 +RNG_SEED: 0 +OUTPUT_DIR: . +LOG_MODEL_INFO: False diff --git a/projects/ssl/README.md b/projects/contrastive_ssl/README.md similarity index 73% rename from projects/ssl/README.md rename to projects/contrastive_ssl/README.md index 4ffae6586..231b44cab 100644 --- a/projects/ssl/README.md +++ b/projects/contrastive_ssl/README.md @@ -13,21 +13,21 @@ In CVPR, 2021. [[Paper](https://arxiv.org/abs/2104.14558)] | method | ρ| architecture | size | frames x sampling | pretrain data | K400-linear | UCF101-split1 | AVA | SSv2 | model | config | | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | -| MoCo | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 66.6 | 91.3 | 19.7 | 52.7 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/MoCo_SlowR50_8x8_T2_epoch_00200.pyth) | ssl/MoCo_SlowR50_8x8 | -| BYOL | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 67.4 | 94.0 | 22.8 | 54.4 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/BYOL_SlowR50_8x8_T2_epoch_00200.pyth) | ssl/BYOL_SlowR50_8x8 | -| SimCLR | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 61.5 | 88.3 | 17.5 | 51.4 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/SimCLR_SlowR50_8x8_T2_epoch_00200.pyth) | ssl/SimCLR_SlowR50_8x8 | -| SwAV | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 62.6 | 90.2 | 19.2 | 52.5| [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/SwAV_SlowR50_8x8_T2_epoch_00200.pyth) | ssl/SwAV_SlowR50_8x8 | -| MoCo | 4 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 71.0 | 94.5 | 21.9 | 54.0 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/MoCo_SlowR50_8x8_T4_epoch_00200.pyth) | ssl/MoCo_SlowR50_8x8 | -| BYOL | 4 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 70.1 | 94.7 | xx | xx | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/BYOL_SlowR50_8x8_T4_epoch_00200.pyth) | ssl/BYOL_SlowR50_8x8 | -| BYOL | 4 | Slow-only | R50 | 16 x 4 | Kinetics-400 | 71.1 | 95.4 | xx | xx | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/BYOL_SlowR50_16x4_T4_epoch_00200.pyth) | ssl/BYOL_SlowR50_8x8 | +| MoCo | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 66.6 | 91.3 | 19.7 | 52.7 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/MoCo_SlowR50_8x8_T2_epoch_00200.pyth) | contrastive_ssl/MoCo_SlowR50_8x8 | +| BYOL | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 67.4 | 94.0 | 22.8 | 54.4 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/BYOL_SlowR50_8x8_T2_epoch_00200.pyth) | contrastive_ssl/BYOL_SlowR50_8x8 | +| SimCLR | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 61.5 | 88.3 | 17.5 | 51.4 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/SimCLR_SlowR50_8x8_T2_epoch_00200.pyth) | contrastive_ssl/SimCLR_SlowR50_8x8 | +| SwAV | 2 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 62.6 | 90.2 | 19.2 | 52.5| [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/SwAV_SlowR50_8x8_T2_epoch_00200.pyth) | contrastive_ssl/SwAV_SlowR50_8x8 | +| MoCo | 4 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 71.0 | 94.5 | 21.9 | 54.0 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/MoCo_SlowR50_8x8_T4_epoch_00200.pyth) | contrastive_ssl/MoCo_SlowR50_8x8 | +| BYOL | 4 | Slow-only | R50 | 8 x 8 | Kinetics-400 | 70.1 | 94.7 | xx | xx | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/BYOL_SlowR50_8x8_T4_epoch_00200.pyth) | contrastive_ssl/BYOL_SlowR50_8x8 | +| BYOL | 4 | Slow-only | R50 | 16 x 4 | Kinetics-400 | 71.1 | 95.4 | xx | xx | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/videomoco_models/BYOL_SlowR50_16x4_T4_epoch_00200.pyth) | contrastive_ssl/BYOL_SlowR50_8x8 | ## Getting started -To use self-supervised learning techniques please refer to the configs under `configs/ssl`, or see the [MODEL_ZOO.md](https://github.com/facebookresearch/SlowFast/blob/master/MODEL_ZOO.md) for pre-trained models. See [paper](https://arxiv.org/abs/2104.14558) for details. For example, the command +To use self-supervised learning techniques please refer to the configs under `configs/contrastive_ssl`, or see the [MODEL_ZOO.md](https://github.com/facebookresearch/SlowFast/blob/master/MODEL_ZOO.md) for pre-trained models. See [paper](https://arxiv.org/abs/2104.14558) for details. For example, the command ``` python tools/run_net.py \ - --cfg configs/Kinetics/ssl/MoCo_SlowR50_8x8.yaml \ + --cfg configs/Kinetics/contrastive_ssl/MoCo_SlowR50_8x8.yaml \ DATA.PATH_TO_DATA_DIR path_to_your_dataset \ ``` @@ -40,6 +40,6 @@ If you find this useful for your research, please consider citing the paper usin @inproceedings{videossl2021, Author = {Christoph Feichtenhofer, Haoqi Fan, Bo Xiong, Ross Girshick, Kaiming He}, Title = {A Large-Scale Study on Unsupervised Spatiotemporal Representation Learning}, - Booktitle = {{CVPR}}, + Booktitle = {CVPR}, Year = {2021}} ``` diff --git a/projects/mae/README.md b/projects/mae/README.md new file mode 100644 index 000000000..c2a3d639e --- /dev/null +++ b/projects/mae/README.md @@ -0,0 +1,53 @@ +# Masked Autoencoders As Spatiotemporal Learners +[Christoph Feichtenhofer*](http://feichtenhofer.github.io/), [Haoqi Fan*](https://haoqifan.github.io/), [Yanghao Li](https://lyttonhao.github.io/), [Kaiming He](http://kaiminghe.com/) +
+Technical report, arXiv, May 2022. [[Paper](https://arxiv.org/abs/2205.09113)] +
+
+ +
+
+ +## Results & Models + +### **Kinetics-400**; configs are under configs/masked_ssl/ + + +| name | frame length x sample rate | top1 | Flops (G) x views | #params (M) | config pre-train (PT) | config fine-tune | model PT | +| ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | +| ViT-B | 16 x 4 | 81.3 | 180 x 3 x 7 | 87 | k400_VIT_B_16x4_MAE_PT | k400_VIT_B_16x4_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/VIT_B_16x4_MAE_PT.pyth) | +| ViT-L | 16 x 4 | 84.8 | 598 x 3 x 7 | 304 | k400_VIT_L_16x4_MAE_PT | k400_VIT_L_16x4_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/VIT_L_16x4_MAE_PT.pyth) | +| ViT-H | 16 x 4 | 85.1 | 1193 x 3 x 7 | 632 | k400_VIT_H_16x4_MAE_PT | k400_VIT_H_16x4_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/VIT_H_16x4_MAE_PT_400e_85.1.pyth) | + + +## Getting started +To use self-supervised learning techniques please refer to the configs under `configs/masked_ssl`. For example, the command + +``` +python tools/run_net.py \ + --cfg configs/masked_ssl/k400_VIT_L_16x4_MAE_PT.yaml \ + DATA.PATH_TO_DATA_DIR path_to_your_Kinetics_dataset +``` + +should train an MAE ViT-L model on the Kinetics-400 dataset, and the command + +``` +python tools/run_net.py \ + --cfg configs/masked_ssl/k400_VIT_L_16x4_FT.yaml \ + DATA.PATH_TO_DATA_DIR path_to_your_Kinetics_dataset \ + TRAIN.CHECKPOINT_FILE_PATH path_to_your_pretrain_checkpoint +``` + +will fine-tune the resulting model, after passing the checkpoint path to the config. + + +## Reference +If you find this useful for your research, please consider citing the paper using the following BibTeX entry. +```BibTeX +@article{feichtenhofer2022masked, + title={Masked Autoencoders As Spatiotemporal Learners}, + author={Feichtenhofer, Christoph and Fan, Haoqi and Li, Yanghao and He, Kaiming}, + journal={arXiv preprint arXiv:2205.09113}, + year={2022} +} +``` diff --git a/projects/maskfeat/README.md b/projects/maskfeat/README.md new file mode 100644 index 000000000..6c6678606 --- /dev/null +++ b/projects/maskfeat/README.md @@ -0,0 +1,81 @@ +# Masked Feature Prediction for Self-Supervised Visual Pre-Training +[Chen Wei*](https://weichen582.github.io/), [Haoqi Fan](https://haoqifan.github.io/), [Saining Xie](https://www.sainingxie.com/), [Chao-Yuan Wu](https://chaoyuan.org/), [Alan Yuille](https://www.cs.jhu.edu/~ayuille/), [Christoph Feichtenhofer*](http://feichtenhofer.github.io/) +
+In CVPR, 2022. [[Paper](https://arxiv.org/abs/2112.09133)] +
+
+ +
+
+ +## Results & Models + +### **ImageNet-1K**; configs are under configs/masked_ssl/ + + +| name | top1 | config pre-train (PT) | config fine-tune | model PT | +| ------------- | ------------- | ------------- | ------------- | ------------- | +| ViT-B | 84.0 | in1k_VIT_B_MaskFeat_PT | in1k_VIT_B_MaskFeat_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/in1k_VIT_B_MaskFeat_PT_epoch_01600.pyth) | +| ViT-L | 85.7 | in1k_VIT_L_MaskFeat_PT | in1k_VIT_L_MaskFeat_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/in1k_VIT_L_MaskFeat_PT_epoch_01600.pyth) | + + + +### **Kinetics-400**; configs are under configs/masked_ssl/ + + +| name | frame length x sample rate | top1 | Flops (G) x views | #params (M) | config pre-train (PT) | config fine-tune | model PT | +| ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | +| MViT-S | 16 x 4 | 82.2 | 71 x 1 x 10 | 36 | k400_MVITv2_S_16x4_MaskFeat_PT | k400_MVITv2_S_16x4_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/k400_MVIT_S_MaskFeat_PT_epoch_00300.pyth) | +| MViT-L | 16 x 4 | 84.3 | 377 x 1 x 10 | 218 | k400_MVITv2_L_16x4_MaskFeat_PT | k400_MVITv2_L_16x4_FT | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/masked_models/k400_MVIT_L_MaskFeat_PT_epoch_00800.pyth) | + + + +## Getting started +To use self-supervised learning techniques please refer to the configs under `configs/masked_ssl`. For example, the command + +``` +python tools/run_net.py \ + --cfg configs/masked_ssl/k400_MVITv2_L_16x4_MaskFeat_PT.yaml \ + DATA.PATH_TO_DATA_DIR path_to_your_Kinetics_dataset +``` + +should train a MaskFeat MViT-L model on the Kinetics-400 dataset, and the command + +``` +python tools/run_net.py \ + --cfg configs/masked_ssl/k400_MVITv2_L_16x4_FT.yaml \ + DATA.PATH_TO_DATA_DIR path_to_your_Kinetics_dataset \ + TRAIN.CHECKPOINT_FILE_PATH path_to_your_pretrain_checkpoint +``` + +will fine-tune the resulting model, after passing the checkpoint path to the config. + +For images, the command + +``` +python tools/run_net.py \ + --cfg configs/masked_ssl/in1k_VIT_B_MaskFeat_PT.yaml \ + DATA.PATH_TO_DATA_DIR path_to_your_ImageNet_dataset +``` + +should train a MaskFeat ViT-B model on the ImageNet dataset, and the command + +``` +python tools/run_net.py \ + --cfg configs/masked_ssl/in1k_VIT_B_FT.yaml \ + DATA.PATH_TO_DATA_DIR path_to_your_ImageNet_dataset \ + TRAIN.CHECKPOINT_FILE_PATH path_to_your_pretrain_checkpoint +``` + +will fine-tune the resulting model, after passing the checkpoint path to the config. + +## Reference +If you find this useful for your research, please consider citing the paper using the following BibTeX entry. +```BibTeX +@InProceedings{wei2022masked, + author = {Wei, Chen and Fan, Haoqi and Xie, Saining and Wu, Chao-Yuan and Yuille, Alan and Feichtenhofer, Christoph}, + title = {Masked Feature Prediction for Self-Supervised Visual Pre-Training}, + booktitle = {CVPR}, + year = {2022}, +} +``` diff --git a/projects/mvitv2/README.md b/projects/mvitv2/README.md index 9e2e0fa0f..6844df3ca 100644 --- a/projects/mvitv2/README.md +++ b/projects/mvitv2/README.md @@ -29,13 +29,34 @@ MViT is a multiscale transformer which serves as a general vision backbone for d | MViTv2-B | 32 x 3 | 82.9 | 95.7 | 225 x 1 x 5 | 51.2 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvitv2/pysf_video_models/MViTv2_B_32x3_k400_f304025456.pyth) | Kinetics/MVITv2_B_32x3 | | MViTv2-L | 40 x 3 | 86.1 | 97.0 | 2828 x 3 x 5 | 217.6 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvitv2/pysf_video_models/MViTv2_L_40x3_k400_f306903192.pyth) | Kinetics/MVITv2_L_40x3_test | + +### SSv2 + + +| name | pretrain | frame length x sample rate | top1 | top5 | Flops (G) x views | #params (M) | model | config | +| ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | ------------- | +| MViTv2-S | K400 | 16 x 4 | 68.2 | 91.4 | 64 x 3 x 1 | 34.4 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvitv2/pysf_video_models/MViTv2_S_16x4_ssv2_f308341823.pyth) | SSv2/MVITv2_S_16x4 | +| MViTv2-B | K400 | 32 x 3 | 70.5 | 92.7 | 225 x 3 x 1 | 51.1 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvitv2/pysf_video_models/MViTv2_B_32x3_ssv2_f305803282.pyth) | SSv2/MVITv2_B_32x3 | +| MViTv2-L | IN21K + K400 | 40 x 3 | 73.3 | 94.1 | 2828 x 3 x 1 | 213.1 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/mvitv2/pysf_video_models/MViTv2_L_40x3_ssv2_f309603968.pyth) | SSv2/MVITv2_L_40x3 | + + +### ImageNet-1K + +| name | resolution |acc@1 | #params | FLOPs | 1k model | +|:---:|:---:|:---:|:---:| :---:|:---:| +| MViTv2-T | 224x224 | 82.3 | 24M | 4.7G | [model](https://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_T_in1k.pyth) | +| MViTv2-S | 224x224 | 83.6 | 35M | 7.0G | [model](https://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_S_in1k.pyth) | +| MViTv2-B | 224x224 | 84.4 | 52M | 10.2G | [model](https://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_B_in1k.pyth) | + +For more ImageNet results, please check the MViTv2 for image classification [repo](https://github.com/facebookresearch/mvit). + ## Get started Here we can train a standard MViTv2 model from scratch by: ``` python tools/run_net.py \ - --cfg configs/Kinetics/MVIT-B.yaml \ + --cfg configs/Kinetics/MVITv2_S_16x4.yaml \ DATA.PATH_TO_DATA_DIR path_to_your_dataset \ ``` diff --git a/projects/rev/README.md b/projects/rev/README.md new file mode 100644 index 000000000..bdc8f0848 --- /dev/null +++ b/projects/rev/README.md @@ -0,0 +1,78 @@ +# Reversible Vision Transformers + +Official PyTorch implementation of **Rev-ViT** and **Rev-MViT** models, from the following paper: +
[Reversible Vision Transformers](https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf) +[Karttikeya Mangalam](https://karttikeya.github.io/)\*, [Haoqi Fan](https://haoqifan.github.io/)\, [Yanghao Li](https://lyttonhao.github.io/)\, [Chao-Yuan Wu](https://chaoyuan.org/)\, [Bo Xiong](https://www.cs.utexas.edu/~bxiong/)\, [Christoph Feichtenhofer](http://feichtenhofer.github.io/)\*, [Jitendra Malik](http://people.eecs.berkeley.edu/~malik/) +CVPR 2022 (Oral) + +**Project Homepage**: https://karttikeya.github.io/publication/revvit/ + +
+
+ +
+
+ +## Pretrained Models + +### ImageNet + +| Architecture | #params (M) | FLOPs (G) | Top1 | weights | Config | +| ------------ | ----------- | --------- | ---- | ------- | --------------------- | +| Rev-ViT-S | 22 | 4.6 | 79.9 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_VIT_S.pyth) | ImageNet/REV_VIT_S | +| Rev-ViT-B | 87 | 17.6 | 81.8 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_VIT_B.pyth) | ImageNet/REV_VIT_B | +| Rev-MViT-B | 39 | 8.7 | 82.9* | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_MVIT_B.pyth) | ImageNet/REV_MVIT_B_16_CONV | + +*improved from 82.5% reported in the [Paper Table 1](https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf). + +### Kinetics 400 + +| Architecture | frame length x sample rate | Top1 | Top5 | Flops (G) x views | #params (M) | weights | config | +| ------------ | -------------------------- | ---- | ---- | ----------------- | ----------- | ------- | --------------------------------- | +| Rev-MViT-B | 16 x 4 | 78.4 | 93.4 | 64 x 1 x 5 | 34.9 | [`link`](https://dl.fbaipublicfiles.com/pyslowfast/rev/REV_MVIT_B_16x4.pyth) | Kinetics/REV_MVIT_B_16x4_CONV | + +## Getting started + +To use Rev-ViT (or Rev-MViT) image models please refer to the configs under `configs/ImageNet`, and see [paper](https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf) for details. For example, the command + +``` +python tools/run_net.py \ + +--cfg configs/ImageNet/REV_VIT_B.yaml \ + +DATA.PATH_TO_DATA_DIR path_to_your_dataset \ +NUM_GPUS 1 \ +``` + +should train and test a Reversible ViT Base Image model (trained with DeiT recipe) on your dataset. Please refer to [general repo-level](https://github.com/facebookresearch/SlowFast/blob/main/GETTING_STARTED.md) instruction for futher details. + +For Rev-MViT video models, please run the `configs/Kinetics` configs as, +``` + +python tools/run_net.py \ + +--cfg configs/Kinetics/REV_MVIT_B_16x4_CONV.yaml \ + +DATA.PATH_TO_DATA_DIR path_to_your_dataset \ +NUM_GPUS 1 \ +``` + +## Citing Rev-ViT \& Rev-MViT + +If you find Reversible Models useful for your research, please consider citing the paper using the following BibTeX entry. + +```BibTeX + + +@inproceedings{mangalam2022, + + title = {Reversible Vision Transformers}, + + author = {Mangalam, Karttikeya and Fan, Haoqi and Li, Yanghao and Wu, Chao-Yuan and Xiong, Bo and Feichtenhofer, Christoph and Malik, Jitendra}, + + booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, + + year = {2022}, +} + +``` diff --git a/projects/rev/teaser.png b/projects/rev/teaser.png new file mode 100644 index 000000000..39abdfd83 Binary files /dev/null and b/projects/rev/teaser.png differ diff --git a/slowfast/config/defaults.py b/slowfast/config/defaults.py index 3be29ea03..2947d821c 100644 --- a/slowfast/config/defaults.py +++ b/slowfast/config/defaults.py @@ -2,7 +2,9 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. """Configs.""" + import math + from fvcore.common.config import CfgNode from . import custom_config @@ -33,7 +35,7 @@ # momentum for momentum encoder updates _C.CONTRASTIVE.MOMENTUM = 0.5 -# wether to anneal momentum to value above with cosine schedule +# whether to anneal momentum to value above with cosine schedule _C.CONTRASTIVE.MOMENTUM_ANNEALING = False # either memorybank, moco, simclr, byol, swav @@ -131,6 +133,10 @@ # If True Train the model, else skip training. _C.TRAIN.ENABLE = True +# Kill training if loss explodes over this ratio from the previous 5 measurements. +# Only enforced if > 0.0 +_C.TRAIN.KILL_LOSS_EXPLOSION_FACTOR = 0.0 + # Dataset. _C.TRAIN.DATASET = "kinetics" @@ -164,6 +170,9 @@ # If True, use FP16 for activations _C.TRAIN.MIXED_PRECISION = False +# if True, inflate some params from imagenet model. +_C.TRAIN.CHECKPOINT_IN_INIT = False + # ---------------------------------------------------------------------------- # # Augmentation options. # ---------------------------------------------------------------------------- # @@ -198,6 +207,32 @@ # Do not random erase first (clean) augmentation split. _C.AUG.RE_SPLIT = False +# Whether to generate input mask during image processing. +_C.AUG.GEN_MASK_LOADER = False + +# If True, masking mode is "tube". Default is "cube". +_C.AUG.MASK_TUBE = False + +# If True, masking mode is "frame". Default is "cube". +_C.AUG.MASK_FRAMES = False + +# The size of generated masks. +_C.AUG.MASK_WINDOW_SIZE = [8, 7, 7] + +# The ratio of masked tokens out of all tokens. Also applies to MViT supervised training +_C.AUG.MASK_RATIO = 0.0 + +# The maximum number of a masked block. None means no maximum limit. (Used only in image MaskFeat.) +_C.AUG.MAX_MASK_PATCHES_PER_BLOCK = None + +# ---------------------------------------------------------------------------- # +# Masked pretraining visualization options. +# ---------------------------------------------------------------------------- # +_C.VIS_MASK = CfgNode() + +# Whether to do visualization. +_C.VIS_MASK.ENABLE = False + # ---------------------------------------------------------------------------- # # MipUp options. # ---------------------------------------------------------------------------- # @@ -250,6 +285,8 @@ _C.TEST.CHECKPOINT_TYPE = "pytorch" # Path to saving prediction results file. _C.TEST.SAVE_RESULTS_PATH = "" + +_C.TEST.NUM_TEMPORAL_CLIPS = [] # ----------------------------------------------------------------------------- # ResNet options # ----------------------------------------------------------------------------- @@ -372,6 +409,7 @@ "slow", "x3d", "mvit", + "maskmvit", ] # Model architectures that has multiple pathways. @@ -444,6 +482,9 @@ # Drop path rate for the tranfomer. _C.MVIT.DROPPATH_RATE = 0.1 +# The initial value of layer scale gamma. Set 0.0 to disable layer scale. +_C.MVIT.LAYER_SCALE_INIT_VALUE = 0.0 + # Depth of the transformer. _C.MVIT.DEPTH = 16 @@ -460,7 +501,7 @@ # Stride size for the Pool KV at layer i. # Format: [[i, stride_t_i, stride_h_i, stride_w_i], ...,] -_C.MVIT.POOL_KV_STRIDE = None +_C.MVIT.POOL_KV_STRIDE = [] # Initial stride size for KV at layer 1. The stride size will be further reduced with # the raio of MVIT.DIM_MUL. If will overwrite MVIT.POOL_KV_STRIDE if not None. @@ -507,6 +548,85 @@ # If True, using separate linear layers for Q, K, V in attention blocks. _C.MVIT.SEPARATE_QKV = False +# The initialization scale factor for the head parameters. +_C.MVIT.HEAD_INIT_SCALE = 1.0 + +# Whether to use the mean pooling of all patch tokens as the output. +_C.MVIT.USE_MEAN_POOLING = False + +# If True, use frozen sin cos positional embedding. +_C.MVIT.USE_FIXED_SINCOS_POS = False + +# ----------------------------------------------------------------------------- +# Masked pretraining options +# ----------------------------------------------------------------------------- +_C.MASK = CfgNode() + +# Whether to enable Masked style pretraining. +_C.MASK.ENABLE = False + +# Whether to enable MAE (discard encoder tokens). +_C.MASK.MAE_ON = False + +# Whether to enable random masking in mae +_C.MASK.MAE_RND_MASK = False + +# Whether to do random masking per-frame in mae +_C.MASK.PER_FRAME_MASKING = False + +# only predict loss on temporal strided patches, or predict full time extent +_C.MASK.TIME_STRIDE_LOSS = True + +# Whether to normalize the pred pixel loss +_C.MASK.NORM_PRED_PIXEL = True + +# Whether to fix initialization with inverse depth of layer for pretraining. +_C.MASK.SCALE_INIT_BY_DEPTH = False + +# Base embedding dimension for the decoder transformer. +_C.MASK.DECODER_EMBED_DIM = 512 + +# Base embedding dimension for the decoder transformer. +_C.MASK.DECODER_SEP_POS_EMBED = False + +# Use a KV kernel in decoder? +_C.MASK.DEC_KV_KERNEL = [] + +# Use a KV stride in decoder? +_C.MASK.DEC_KV_STRIDE = [] + +# The depths of features which are inputs of the prediction head. +_C.MASK.PRETRAIN_DEPTH = [15] + +# The type of Masked pretraining prediction head. +# Can be "separate", "separate_xformer". +_C.MASK.HEAD_TYPE = "separate" + +# The depth of MAE's decoder +_C.MASK.DECODER_DEPTH = 0 + +# The weight of HOG target loss. +_C.MASK.PRED_HOG = False +# Reversible Configs +_C.MVIT.REV = CfgNode() + +# Enable Reversible Model +_C.MVIT.REV.ENABLE = False + +# Method to fuse the reversible paths +# see :class: `TwoStreamFusion` for all the options +_C.MVIT.REV.RESPATH_FUSE = "concat" + +# Layers to buffer activations at +# (at least Q-pooling layers needed) +_C.MVIT.REV.BUFFER_LAYERS = [] + +# 'conv' or 'max' operator for the respath in Qpooling +_C.MVIT.REV.RES_PATH = "conv" + +# Method to merge hidden states before Qpoolinglayers +_C.MVIT.REV.PRE_Q_FUSION = "avg" + # ----------------------------------------------------------------------------- # SlowFast options # ----------------------------------------------------------------------------- @@ -605,6 +725,9 @@ # Decoding backend, options include `pyav` or `torchvision` _C.DATA.DECODING_BACKEND = "torchvision" +# Decoding resize to short size (set to native size for best speed) +_C.DATA.DECODING_SHORT_SIZE = 256 + # if True, sample uniformly in [1 / max_scale, 1 / min_scale] and take a # reciprocal to get the scale. If False, take a uniform sample from # [min_scale, max_scale]. @@ -668,6 +791,18 @@ _C.DATA.SSL_BLUR_SIGMA_MAX = [0.0, 2.0] +# If combine train/val split as training for in21k +_C.DATA.IN22K_TRAINVAL = False + +# If not None, use IN1k as val split when training in21k +_C.DATA.IN22k_VAL_IN1K = "" + +# Large resolution models may use different crop ratios +_C.DATA.IN_VAL_CROP_RATIO = 0.875 # 224/256 = 0.875 + +# don't use real video for kinetics.py +_C.DATA.DUMMY_LOAD = False + # ---------------------------------------------------------------------------- # # Optimizer options # ---------------------------------------------------------------------------- # @@ -739,6 +874,11 @@ # LARS optimizer _C.SOLVER.LARS_ON = False +# The layer-wise decay of learning rate. Set to 1. to disable. +_C.SOLVER.LAYER_DECAY = 1.0 + +# Adam's beta +_C.SOLVER.BETAS = (0.9, 0.999) # ---------------------------------------------------------------------------- # # Misc options # ---------------------------------------------------------------------------- # diff --git a/slowfast/datasets/__init__.py b/slowfast/datasets/__init__.py index ae97afdae..9e14e0a51 100644 --- a/slowfast/datasets/__init__.py +++ b/slowfast/datasets/__init__.py @@ -2,7 +2,7 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. from .ava_dataset import Ava # noqa -from .build import DATASET_REGISTRY, build_dataset # noqa +from .build import build_dataset, DATASET_REGISTRY # noqa from .charades import Charades # noqa from .imagenet import Imagenet # noqa from .kinetics import Kinetics # noqa diff --git a/slowfast/datasets/ava_dataset.py b/slowfast/datasets/ava_dataset.py index a0459599d..280381d27 100644 --- a/slowfast/datasets/ava_dataset.py +++ b/slowfast/datasets/ava_dataset.py @@ -2,13 +2,16 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import logging + import numpy as np import torch -from . import ava_helper as ava_helper -from . import cv2_transform as cv2_transform -from . import transform as transform -from . import utils as utils +from . import ( + ava_helper as ava_helper, + cv2_transform as cv2_transform, + transform as transform, + utils as utils, +) from .build import DATASET_REGISTRY logger = logging.getLogger(__name__) @@ -60,9 +63,7 @@ def _load_data(self, cfg): ) = ava_helper.load_image_lists(cfg, is_train=(self._split == "train")) # Loading annotations for boxes and labels. - boxes_and_labels = ava_helper.load_boxes_and_labels( - cfg, mode=self._split - ) + boxes_and_labels = ava_helper.load_boxes_and_labels(cfg, mode=self._split) assert len(boxes_and_labels) == len(self._image_paths) @@ -155,9 +156,7 @@ def _images_and_boxes_preprocessing_cv2(self, imgs, boxes): # Short side to test_scale. Non-local and STRG uses 256. imgs = [cv2_transform.scale(self._crop_size, img) for img in imgs] boxes = [ - cv2_transform.scale_boxes( - self._crop_size, boxes[0], height, width - ) + cv2_transform.scale_boxes(self._crop_size, boxes[0], height, width) ] imgs, boxes = cv2_transform.spatial_shift_crop_list( self._crop_size, imgs, 1, boxes=boxes @@ -171,9 +170,7 @@ def _images_and_boxes_preprocessing_cv2(self, imgs, boxes): # Short side to test_scale. Non-local and STRG uses 256. imgs = [cv2_transform.scale(self._crop_size, img) for img in imgs] boxes = [ - cv2_transform.scale_boxes( - self._crop_size, boxes[0], height, width - ) + cv2_transform.scale_boxes(self._crop_size, boxes[0], height, width) ] if self._test_force_flip: @@ -181,9 +178,7 @@ def _images_and_boxes_preprocessing_cv2(self, imgs, boxes): 1, imgs, order="HWC", boxes=boxes ) else: - raise NotImplementedError( - "Unsupported split mode {}".format(self._split) - ) + raise NotImplementedError("Unsupported split mode {}".format(self._split)) # Convert image to CHW keeping BGR order. imgs = [cv2_transform.HWC2CHW(img) for img in imgs] @@ -227,9 +222,7 @@ def _images_and_boxes_preprocessing_cv2(self, imgs, boxes): ] # Concat list of images to single ndarray. - imgs = np.concatenate( - [np.expand_dims(img, axis=1) for img in imgs], axis=1 - ) + imgs = np.concatenate([np.expand_dims(img, axis=1) for img in imgs], axis=1) if not self._use_bgr: # Convert image format from BGR to RGB. @@ -274,9 +267,7 @@ def _images_and_boxes_preprocessing(self, imgs, boxes): max_size=self._jitter_max_scale, boxes=boxes, ) - imgs, boxes = transform.random_crop( - imgs, self._crop_size, boxes=boxes - ) + imgs, boxes = transform.random_crop(imgs, self._crop_size, boxes=boxes) # Random flip. imgs, boxes = transform.horizontal_flip(0.5, imgs, boxes=boxes) @@ -310,9 +301,7 @@ def _images_and_boxes_preprocessing(self, imgs, boxes): if self._test_force_flip: imgs, boxes = transform.horizontal_flip(1, imgs, boxes=boxes) else: - raise NotImplementedError( - "{} split not supported yet!".format(self._split) - ) + raise NotImplementedError("{} split not supported yet!".format(self._split)) # Do color augmentation (after divided by 255.0). if self._split == "train" and self._use_color_augmentation: @@ -343,9 +332,7 @@ def _images_and_boxes_preprocessing(self, imgs, boxes): # Note that Kinetics pre-training uses RGB! imgs = imgs[:, [2, 1, 0], ...] - boxes = transform.clip_boxes_to_image( - boxes, self._crop_size, self._crop_size - ) + boxes = transform.clip_boxes_to_image(boxes, self._crop_size, self._crop_size) return imgs, boxes @@ -403,16 +390,12 @@ def __getitem__(self, idx): # T H W C -> T C H W. imgs = imgs.permute(0, 3, 1, 2) # Preprocess images and boxes. - imgs, boxes = self._images_and_boxes_preprocessing( - imgs, boxes=boxes - ) + imgs, boxes = self._images_and_boxes_preprocessing(imgs, boxes=boxes) # T C H W -> C T H W. imgs = imgs.permute(1, 0, 2, 3) else: # Preprocess images and boxes - imgs, boxes = self._images_and_boxes_preprocessing_cv2( - imgs, boxes=boxes - ) + imgs, boxes = self._images_and_boxes_preprocessing_cv2(imgs, boxes=boxes) # Construct label arrays. label_arrs = np.zeros((len(labels), self._num_classes), dtype=np.int32) diff --git a/slowfast/datasets/ava_helper.py b/slowfast/datasets/ava_helper.py index 7a75a7f65..79f4f3d68 100644 --- a/slowfast/datasets/ava_helper.py +++ b/slowfast/datasets/ava_helper.py @@ -29,9 +29,7 @@ def load_image_lists(cfg, is_train): """ list_filenames = [ os.path.join(cfg.AVA.FRAME_LIST_DIR, filename) - for filename in ( - cfg.AVA.TRAIN_LISTS if is_train else cfg.AVA.TEST_LISTS - ) + for filename in (cfg.AVA.TRAIN_LISTS if is_train else cfg.AVA.TEST_LISTS) ] image_paths = defaultdict(list) video_name_to_idx = {} @@ -53,15 +51,11 @@ def load_image_lists(cfg, is_train): data_key = video_name_to_idx[video_name] - image_paths[data_key].append( - os.path.join(cfg.AVA.FRAME_DIR, row[3]) - ) + image_paths[data_key].append(os.path.join(cfg.AVA.FRAME_DIR, row[3])) image_paths = [image_paths[i] for i in range(len(image_paths))] - logger.info( - "Finished loading image paths from: %s" % ", ".join(list_filenames) - ) + logger.info("Finished loading image paths from: %s" % ", ".join(list_filenames)) return image_paths, video_idx_to_name @@ -95,18 +89,14 @@ def load_boxes_and_labels(cfg, mode): detect_thresh = cfg.AVA.DETECTION_SCORE_THRESH # Only select frame_sec % 4 = 0 samples for validation if not # set FULL_TEST_ON_VAL. - boxes_sample_rate = ( - 4 if mode == "val" and not cfg.AVA.FULL_TEST_ON_VAL else 1 - ) + boxes_sample_rate = 4 if mode == "val" and not cfg.AVA.FULL_TEST_ON_VAL else 1 all_boxes, count, unique_box_count = parse_bboxes_file( ann_filenames=ann_filenames, ann_is_gt_box=ann_is_gt_box, detect_thresh=detect_thresh, boxes_sample_rate=boxes_sample_rate, ) - logger.info( - "Finished loading annotations from: %s" % ", ".join(ann_filenames) - ) + logger.info("Finished loading annotations from: %s" % ", ".join(ann_filenames)) logger.info("Detection threshold: {}".format(detect_thresh)) logger.info("Number of unique boxes: %d" % unique_box_count) logger.info("Number of annotations: %d" % count) @@ -147,9 +137,7 @@ def sec_to_frame(sec): continue if len(boxes_and_labels[video_idx][sec]) > 0: - keyframe_indices.append( - (video_idx, sec_idx, sec, sec_to_frame(sec)) - ) + keyframe_indices.append((video_idx, sec_idx, sec, sec_to_frame(sec))) keyframe_boxes_and_labels[video_idx].append( boxes_and_labels[video_idx][sec] ) @@ -179,9 +167,7 @@ def get_num_boxes_used(keyframe_indices, keyframe_boxes_and_labels): return count -def parse_bboxes_file( - ann_filenames, ann_is_gt_box, detect_thresh, boxes_sample_rate=1 -): +def parse_bboxes_file(ann_filenames, ann_is_gt_box, detect_thresh, boxes_sample_rate=1): """ Parse AVA bounding boxes files. Args: diff --git a/slowfast/datasets/charades.py b/slowfast/datasets/charades.py index 68a943dde..07814e9d6 100644 --- a/slowfast/datasets/charades.py +++ b/slowfast/datasets/charades.py @@ -4,10 +4,10 @@ import os import random from itertools import chain as chain -import torch -import torch.utils.data import slowfast.utils.logging as logging +import torch +import torch.utils.data from slowfast.utils.env import pathmgr from . import utils as utils @@ -64,9 +64,7 @@ def __init__(self, cfg, mode, num_retries=10): if self.mode in ["train", "val"]: self._num_clips = 1 elif self.mode in ["test"]: - self._num_clips = ( - cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS - ) + self._num_clips = cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS logger.info("Constructing Charades {}...".format(mode)) self._construct_loader() @@ -79,9 +77,7 @@ def _construct_loader(self): self.cfg.DATA.PATH_TO_DATA_DIR, "{}.csv".format("train" if self.mode == "train" else "val"), ) - assert pathmgr.exists(path_to_file), "{} dir not found".format( - path_to_file - ) + assert pathmgr.exists(path_to_file), "{} dir not found".format(path_to_file) (self._path_to_videos, self._labels) = utils.load_image_lists( path_to_file, self.cfg.DATA.PATH_PREFIX, return_list=True ) @@ -91,9 +87,7 @@ def _construct_loader(self): self._labels = utils.convert_to_video_level_labels(self._labels) self._path_to_videos = list( - chain.from_iterable( - [[x] * self._num_clips for x in self._path_to_videos] - ) + chain.from_iterable([[x] * self._num_clips for x in self._path_to_videos]) ) self._labels = list( chain.from_iterable([[x] * self._num_clips for x in self._labels]) @@ -121,8 +115,7 @@ def get_seq_frames(self, index): temporal_sample_index = ( -1 if self.mode in ["train", "val"] - else self._spatial_temporal_idx[index] - // self.cfg.TEST.NUM_SPATIAL_CROPS + else self._spatial_temporal_idx[index] // self.cfg.TEST.NUM_SPATIAL_CROPS ) num_frames = self.cfg.DATA.NUM_FRAMES sampling_rate = utils.get_random_sampling_rate( @@ -189,28 +182,21 @@ def __getitem__(self, index): # Decreasing the scale is equivalent to using a larger "span" # in a sampling grid. min_scale = int( - round( - float(min_scale) - * crop_size - / self.cfg.MULTIGRID.DEFAULT_S - ) + round(float(min_scale) * crop_size / self.cfg.MULTIGRID.DEFAULT_S) ) elif self.mode in ["test"]: # spatial_sample_index is in [0, 1, 2]. Corresponding to left, # center, or right if width is larger than height, and top, middle, # or bottom if height is larger than width. spatial_sample_index = ( - self._spatial_temporal_idx[index] - % self.cfg.TEST.NUM_SPATIAL_CROPS + self._spatial_temporal_idx[index] % self.cfg.TEST.NUM_SPATIAL_CROPS ) min_scale, max_scale, crop_size = [self.cfg.DATA.TEST_CROP_SIZE] * 3 # The testing is deterministic and no jitter should be performed. # min_scale, max_scale, and crop_size are expect to be the same. assert len({min_scale, max_scale, crop_size}) == 1 else: - raise NotImplementedError( - "Does not support {} mode".format(self.mode) - ) + raise NotImplementedError("Does not support {} mode".format(self.mode)) seq = self.get_seq_frames(index) frames = torch.as_tensor( @@ -228,9 +214,7 @@ def __getitem__(self, index): ) # Perform color normalization. - frames = utils.tensor_normalize( - frames, self.cfg.DATA.MEAN, self.cfg.DATA.STD - ) + frames = utils.tensor_normalize(frames, self.cfg.DATA.MEAN, self.cfg.DATA.STD) # T H W C -> C T H W. frames = frames.permute(3, 0, 1, 2) # Perform data augmentation. diff --git a/slowfast/datasets/cv2_transform.py b/slowfast/datasets/cv2_transform.py index e375e50b4..d080f4570 100644 --- a/slowfast/datasets/cv2_transform.py +++ b/slowfast/datasets/cv2_transform.py @@ -2,8 +2,9 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import math -import numpy as np + import cv2 +import numpy as np def clip_boxes_to_image(boxes, height, width): @@ -17,12 +18,8 @@ def clip_boxes_to_image(boxes, height, width): Returns: boxes (ndarray): cropped bounding boxes. """ - boxes[:, [0, 2]] = np.minimum( - width - 1.0, np.maximum(0.0, boxes[:, [0, 2]]) - ) - boxes[:, [1, 3]] = np.minimum( - height - 1.0, np.maximum(0.0, boxes[:, [1, 3]]) - ) + boxes[:, [0, 2]] = np.minimum(width - 1.0, np.maximum(0.0, boxes[:, [0, 2]])) + boxes[:, [1, 3]] = np.minimum(height - 1.0, np.maximum(0.0, boxes[:, [1, 3]])) return boxes @@ -47,18 +44,14 @@ def random_short_side_scale_jitter_list(images, min_size, max_size, boxes=None): height = images[0].shape[0] width = images[0].shape[1] - if (width <= height and width == size) or ( - height <= width and height == size - ): + if (width <= height and width == size) or (height <= width and height == size): return images, boxes new_width = size new_height = size if width < height: new_height = int(math.floor((float(height) / width) * size)) if boxes is not None: - boxes = [ - proposal * float(new_height) / height for proposal in boxes - ] + boxes = [proposal * float(new_height) / height for proposal in boxes] else: new_width = int(math.floor((float(width) / height) * size)) if boxes is not None: @@ -87,9 +80,7 @@ def scale(size, image): """ height = image.shape[0] width = image.shape[1] - if (width <= height and width == size) or ( - height <= width and height == size - ): + if (width <= height and width == size) or (height <= width and height == size): return image new_width = size new_height = size @@ -97,9 +88,7 @@ def scale(size, image): new_height = int(math.floor((float(height) / width) * size)) else: new_width = int(math.floor((float(width) / height) * size)) - img = cv2.resize( - image, (new_width, new_height), interpolation=cv2.INTER_LINEAR - ) + img = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR) return img.astype(np.float32) @@ -115,9 +104,7 @@ def scale_boxes(size, boxes, height, width): Returns: boxes (ndarray): scaled bounding boxes. """ - if (width <= height and width == size) or ( - height <= width and height == size - ): + if (width <= height and width == size) or (height <= width and height == size): return boxes new_width = size @@ -237,9 +224,7 @@ def HWC2CHW(image): return image.transpose([2, 0, 1]) -def color_jitter_list( - images, img_brightness=0, img_contrast=0, img_saturation=0 -): +def color_jitter_list(images, img_brightness=0, img_contrast=0, img_saturation=0): """ Perform color jitter on the list of images. Args: @@ -410,8 +395,7 @@ def random_crop_list(images, size, pad_size=0, order="CHW", boxes=None): # explicitly dealing processing per image order to avoid flipping images. if pad_size > 0: images = [ - pad_image(pad_size=pad_size, image=image, order=order) - for image in images + pad_image(pad_size=pad_size, image=image, order=order) for image in images ] # image format should be CHW. @@ -485,9 +469,7 @@ def random_scale_jitter(image, min_size, max_size): Returns: image (array): scaled image. """ - img_scale = int( - round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size)) - ) + img_scale = int(round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size))) image = scale(img_scale, image) return image @@ -504,9 +486,7 @@ def random_scale_jitter_list(images, min_size, max_size): Returns: images (list): list of scaled image. """ - img_scale = int( - round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size)) - ) + img_scale = int(round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size))) return [scale(img_scale, image) for image in images] @@ -543,12 +523,8 @@ def random_sized_crop(image, size, area_frac=0.08): y_offset = int(y_offset) x_offset = int(x_offset) cropped = image[y_offset : y_offset + h, x_offset : x_offset + w, :] - assert ( - cropped.shape[0] == h and cropped.shape[1] == w - ), "Wrong crop size" - cropped = cv2.resize( - cropped, (size, size), interpolation=cv2.INTER_LINEAR - ) + assert cropped.shape[0] == h and cropped.shape[1] == w, "Wrong crop size" + cropped = cv2.resize(cropped, (size, size), interpolation=cv2.INTER_LINEAR) return cropped.astype(np.float32) return center_crop(size, scale(size, image)) @@ -614,12 +590,10 @@ def random_sized_crop_list(images, size, crop_area_fraction=0.08): croppsed_images = [] for image in images: - cropped = image[ - y_offset : y_offset + h, x_offset : x_offset + w, : - ] - assert ( - cropped.shape[0] == h and cropped.shape[1] == w - ), "Wrong crop size" + cropped = image[y_offset : y_offset + h, x_offset : x_offset + w, :] + assert cropped.shape[0] == h and cropped.shape[1] == w, ( + "Wrong crop size" + ) cropped = cv2.resize( cropped, (size, size), interpolation=cv2.INTER_LINEAR ) diff --git a/slowfast/datasets/decoder.py b/slowfast/datasets/decoder.py index 7ebf713e3..c4f4f1cda 100644 --- a/slowfast/datasets/decoder.py +++ b/slowfast/datasets/decoder.py @@ -3,8 +3,9 @@ import logging import math -import numpy as np import random + +import numpy as np import torch import torchvision.io as io @@ -66,9 +67,7 @@ def get_start_end_idx( start_idx = math.floor(delta / 2) else: # Uniformly sample the clip with the given index. - start_idx = clip_idx * math.floor( - delta / (num_clips_uniform - 1) - ) + start_idx = clip_idx * math.floor(delta / (num_clips_uniform - 1)) else: # Uniformly sample the clip with the given index. start_idx = delta * clip_idx / num_clips_uniform @@ -84,7 +83,7 @@ def get_multiple_start_end_idx( num_clips_uniform, min_delta=0, max_delta=math.inf, - use_offset=False + use_offset=False, ): """ Sample a clip of size clip_size from a video of size video_size and @@ -114,7 +113,7 @@ def sample_clips( min_delta=0, max_delta=math.inf, num_retries=100, - use_offset=False + use_offset=False, ): se_inds = np.empty((0, 2)) dt = np.empty((0)) @@ -125,13 +124,15 @@ def sample_clips( if clip_idx == -1: # Random temporal sampling. start_idx = random.uniform(0, max_start) - else: # Uniformly sample the clip with the given index. + else: # Uniformly sample the clip with the given index. if use_offset: if num_clips_uniform == 1: # Take the center clip if num_clips is 1. start_idx = math.floor(max_start / 2) else: - start_idx = clip_idx * math.floor(max_start / (num_clips_uniform - 1)) + start_idx = clip_idx * math.floor( + max_start / (num_clips_uniform - 1) + ) else: start_idx = max_start * clip_idx / num_clips_uniform @@ -302,10 +303,7 @@ def torchvision_decode( decode_all_video = False # try selective decoding clip_sizes = [ - np.maximum( - 1.0, - sampling_rate[i] * num_frames[i] / target_fps * fps - ) + np.maximum(1.0, sampling_rate[i] * num_frames[i] / target_fps * fps) for i in range(len(sampling_rate)) ] start_end_delta_time = get_multiple_start_end_idx( @@ -319,9 +317,7 @@ def torchvision_decode( ) frames_out = [None] * len(num_frames) for k in range(len(num_frames)): - pts_per_frame = ( - video_meta["video_denominator"] / video_meta["video_fps"] - ) + pts_per_frame = video_meta["video_denominator"] / video_meta["video_fps"] video_start_pts = int(start_end_delta_time[k, 0] * pts_per_frame) video_end_pts = int(start_end_delta_time[k, 1] * pts_per_frame) @@ -505,14 +501,12 @@ def decode( sampling_rate = [sampling_rate[i] for i in ind_clips] num_frames = [num_frames[i] for i in ind_clips] else: - ind_clips = np.arange( - num_decode - ) # clips come temporally ordered from decoder + ind_clips = np.arange(num_decode) # clips come temporally ordered from decoder try: if backend == "pyav": - assert ( - min_delta == -math.inf and max_delta == math.inf - ), "delta sampling not supported in pyav" + assert min_delta == -math.inf and max_delta == math.inf, ( + "delta sampling not supported in pyav" + ) frames_decoded, fps, decode_all_video = pyav_decode( container, sampling_rate, @@ -543,9 +537,7 @@ def decode( max_delta=max_delta, ) else: - raise NotImplementedError( - "Unknown decoding backend {}".format(backend) - ) + raise NotImplementedError("Unknown decoding backend {}".format(backend)) except Exception as e: print("Failed to decode by {} with exception: {}".format(backend, e)) return None, None, None @@ -558,12 +550,9 @@ def decode( frames_decoded = [frames_decoded] num_decoded = len(frames_decoded) clip_sizes = [ - np.maximum( - 1.0, - sampling_rate[i] * num_frames[i] / target_fps * fps - ) - for i in range(len(sampling_rate)) - ] + np.maximum(1.0, sampling_rate[i] * num_frames[i] / target_fps * fps) + for i in range(len(sampling_rate)) + ] if decode_all_video: # full video was decoded (not trimmed yet) assert num_decoded == 1 and start_end_delta_time is None @@ -621,8 +610,7 @@ def decode( start_end_delta_time = start_end_delta_time_ time_diff_aug = time_diff_aug_ assert all( - frames_out[i].shape[0] == num_frames_orig[i] - for i in range(num_decode) + frames_out[i].shape[0] == num_frames_orig[i] for i in range(num_decode) ) return frames_out, start_end_delta_time, time_diff_aug diff --git a/slowfast/datasets/imagenet.py b/slowfast/datasets/imagenet.py index 1a706621d..daf49d313 100644 --- a/slowfast/datasets/imagenet.py +++ b/slowfast/datasets/imagenet.py @@ -1,23 +1,24 @@ # Copyright (c) Facebook, Inc. and its affiliates. import json -import numpy as np import os import random import re -import torch -import torch.utils.data -from PIL import Image -from torchvision import transforms as transforms_tv +import numpy as np import slowfast.datasets.transform as transform import slowfast.utils.logging as logging +import torch +import torch.utils.data +from PIL import Image +from slowfast.models.utils import calc_mvit_feature_geometry # import cv2 from slowfast.utils.env import pathmgr +from torchvision import transforms as transforms_tv from .build import DATASET_REGISTRY -from .transform import transforms_imagenet_train +from .transform import MaskingGenerator, transforms_imagenet_train logger = logging.get_logger(__name__) @@ -41,6 +42,9 @@ def __init__(self, cfg, mode, num_retries=10): self._construct_imdb() else: self._load_imdb() + self.num_videos = len(self._imdb) + self.feat_size, self.feat_stride = calc_mvit_feature_geometry(cfg) + self.dummy_output = None def _load_imdb(self): split_path = os.path.join( @@ -57,9 +61,7 @@ def _construct_imdb(self): logger.info("{} data path: {}".format(self.mode, split_path)) # Images are stored per class in subdirs (format: n) split_files = pathmgr.ls(split_path) - self._class_ids = sorted( - f for f in split_files if re.match(r"^n[0-9]+$", f) - ) + self._class_ids = sorted(f for f in split_files if re.match(r"^n[0-9]+$", f)) # Map ImageNet class ids to contiguous ids self._class_id_cont_id = {v: i for i, v in enumerate(self._class_ids)} # Construct the image db @@ -113,11 +115,7 @@ def _prepare_im_res(self, im_path): im, test_size, spatial_idx=1, scale_size=train_size ) # For training and testing use color normalization - im = transform.color_normalization( - im, self.cfg.DATA.MEAN, self.cfg.DATA.STD - ) - # Convert HWC/RGB/float to CHW/BGR/float format - # im = np.ascontiguousarray(im[:, :, ::-1].transpose([2, 0, 1])) + im = transform.color_normalization(im, self.cfg.DATA.MEAN, self.cfg.DATA.STD) return im def _prepare_im_tf(self, im_path): @@ -144,28 +142,81 @@ def _prepare_im_tf(self, im_path): ) else: t = [] - size = int((256 / 224) * test_size) - t.append( - transforms_tv.Resize( - size, interpolation=3 - ), # to maintain same ratio w.r.t. 224 images - ) - t.append(transforms_tv.CenterCrop(test_size)) + if self.cfg.DATA.IN_VAL_CROP_RATIO == 0.0: + t.append( + transforms_tv.Resize( + (test_size, test_size), + interpolation=transforms_tv.InterpolationMode.BICUBIC, + ), + ) + else: + size = int( + (1.0 / self.cfg.DATA.IN_VAL_CROP_RATIO) * test_size + ) # = 1/0.875 * test_size + t.append( + transforms_tv.Resize( + size, interpolation=transforms_tv.InterpolationMode.BICUBIC + ), # to maintain same ratio w.r.t. 224 images + ) + t.append(transforms_tv.CenterCrop(test_size)) t.append(transforms_tv.ToTensor()) - t.append( - transforms_tv.Normalize(self.cfg.DATA.MEAN, self.cfg.DATA.STD) - ) + t.append(transforms_tv.Normalize(self.cfg.DATA.MEAN, self.cfg.DATA.STD)) aug_transform = transforms_tv.Compose(t) im = aug_transform(im) return im + def _prepare_im_masked(self, im_path): + with pathmgr.open(im_path, "rb") as f: + with Image.open(f) as im: + im = im.convert("RGB") + + if self.mode in ["train", "val"]: + depth = self.cfg.MASK.PRETRAIN_DEPTH[-1] + assert depth == max(self.cfg.MASK.PRETRAIN_DEPTH) + max_mask = self.cfg.AUG.MAX_MASK_PATCHES_PER_BLOCK + # use feat geometry for determining num masks + mask_window_size = self.feat_size[depth][-1] + num_mask = round( + self.feat_size[depth][-1] + * self.feat_size[depth][-2] + * self.cfg.AUG.MASK_RATIO + ) + min_mask = num_mask // 5 + + train_size = self.cfg.DATA.TRAIN_CROP_SIZE + mask_generator = MaskingGenerator( + mask_window_size, + num_masking_patches=num_mask, + max_num_patches=max_mask, + min_num_patches=min_mask, + ) + aug_transform = transforms_imagenet_train( + img_size=(train_size, train_size), + scale=self.cfg.DATA.TRAIN_JITTER_SCALES_RELATIVE, + ratio=self.cfg.DATA.TRAIN_JITTER_ASPECT_RELATIVE, + interpolation=self.cfg.AUG.INTERPOLATION, + color_jitter=self.cfg.AUG.COLOR_JITTER, + auto_augment=self.cfg.AUG.AA_TYPE, + re_prob=0.0, + mean=self.cfg.DATA.MEAN, + std=self.cfg.DATA.STD, + ) + im = aug_transform(im) + mask = mask_generator() + return [im, torch.Tensor(), mask] + else: + raise NotImplementedError + return aug_transform(im) + def __load__(self, index): try: # Load the image im_path = self._imdb[index]["im_path"] # Prepare the image for training / testing if self.cfg.AUG.ENABLE: - if self.mode == "train" and self.cfg.AUG.NUM_SAMPLE > 1: + if self.cfg.AUG.GEN_MASK_LOADER: + return self._prepare_im_masked(im_path) + elif self.mode == "train" and self.cfg.AUG.NUM_SAMPLE > 1: im = [] for _ in range(self.cfg.AUG.NUM_SAMPLE): crop = self._prepare_im_tf(im_path) @@ -181,23 +232,37 @@ def __load__(self, index): return None def __getitem__(self, index): + if self.dummy_output is not None: + return self.dummy_output # if the current image is corrupted, load a different image. for _ in range(self.num_retries): im = self.__load__(index) # Data corrupted, retry with a different image. if im is None: + assert self.mode == "train", f"{index} failed loading" + print(f"{index} failed. retry") index = random.randint(0, len(self._imdb) - 1) else: break # Retrieve the label label = self._imdb[index]["class"] if isinstance(im, list): - label = [label for _ in range(len(im))] - dummy = [torch.Tensor() for _ in range(len(im))] - return im, label, dummy, dummy, {} + if self.cfg.AUG.GEN_MASK_LOADER: + dummy = torch.Tensor() + label = torch.Tensor() + else: + label = [label for _ in range(len(im))] + dummy = [torch.Tensor() for _ in range(len(im))] + if self.cfg.DATA.DUMMY_LOAD: + if self.dummy_output is None: + self.dummy_output = (im, label, index, dummy, {}) + return im, label, index, dummy, {} else: dummy = torch.Tensor() - return [im], label, dummy, dummy, {} + if self.cfg.DATA.DUMMY_LOAD: + if self.dummy_output is None: + self.dummy_output = ([im], label, index, dummy, {}) + return [im], label, index, dummy, {} def __len__(self): return len(self._imdb) diff --git a/slowfast/datasets/kinetics.py b/slowfast/datasets/kinetics.py index 0e5eada55..5411597d6 100644 --- a/slowfast/datasets/kinetics.py +++ b/slowfast/datasets/kinetics.py @@ -1,24 +1,26 @@ #!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -import numpy as np import os import random + +import numpy as np import pandas +import slowfast.utils.logging as logging import torch import torch.utils.data -from torchvision import transforms - -import slowfast.utils.logging as logging from slowfast.utils.env import pathmgr +from torchvision import transforms -from . import decoder as decoder -from . import transform as transform -from . import utils as utils -from . import video_container as container +from . import ( + decoder as decoder, + transform as transform, + utils as utils, + video_container as container, +) from .build import DATASET_REGISTRY from .random_erasing import RandomErasing -from .transform import create_random_augment +from .transform import create_random_augment, MaskingGenerator, MaskingGenerator3D logger = logging.get_logger(__name__) @@ -74,6 +76,7 @@ def __init__(self, cfg, mode, num_retries=100): if self.mode in ["train"] and self.cfg.DATA.LOADER_CHUNK_SIZE > 0 else False ) + self.dummy_output = None # For training or validation mode, one single clip is sampled from every # video. For testing, NUM_ENSEMBLE_VIEWS clips are sampled from every # video. For every clip, NUM_SPATIAL_CROPS is cropped spatially from @@ -81,19 +84,18 @@ def __init__(self, cfg, mode, num_retries=100): if self.mode in ["train", "val"]: self._num_clips = 1 elif self.mode in ["test"]: - self._num_clips = ( - cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS - ) + self._num_clips = cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS logger.info("Constructing Kinetics {}...".format(mode)) self._construct_loader() - self.randaug = False + self.aug = False self.rand_erase = False self.use_temporal_gradient = False self.temporal_gradient_rate = 0.0 + self.cur_epoch = 0 if self.mode == "train" and self.cfg.AUG.ENABLE: - self.randaug = True + self.aug = True if self.cfg.AUG.RE_PROB > 0: self.rand_erase = True @@ -104,9 +106,7 @@ def _construct_loader(self): path_to_file = os.path.join( self.cfg.DATA.PATH_TO_DATA_DIR, "{}.csv".format(self.mode) ) - assert pathmgr.exists(path_to_file), "{} dir not found".format( - path_to_file - ) + assert pathmgr.exists(path_to_file), "{} dir not found".format(path_to_file) self._path_to_videos = [] self._labels = [] @@ -122,9 +122,7 @@ def _construct_loader(self): else: rows = f.read().splitlines() for clip_idx, path_label in enumerate(rows): - fetch_info = path_label.split( - self.cfg.DATA.PATH_LABEL_SEPARATOR - ) + fetch_info = path_label.split(self.cfg.DATA.PATH_LABEL_SEPARATOR) if len(fetch_info) == 2: path, label = fetch_info elif len(fetch_info) == 3: @@ -144,10 +142,10 @@ def _construct_loader(self): self._labels.append(int(label)) self._spatial_temporal_idx.append(idx) self._video_meta[clip_idx * self._num_clips + idx] = {} - assert ( - len(self._path_to_videos) > 0 - ), "Failed to load Kinetics split {} from {}".format( - self._split_idx, path_to_file + assert len(self._path_to_videos) > 0, ( + "Failed to load Kinetics split {} from {}".format( + self._split_idx, path_to_file + ) ) logger.info( "Constructing kinetics dataloader (size: {} skip_rows {}) from {} ".format( @@ -160,12 +158,13 @@ def _set_epoch_num(self, epoch): def _get_chunk(self, path_to_file, chunksize): try: - for chunk in pandas.read_csv( - path_to_file, - chunksize=self.cfg.DATA.LOADER_CHUNK_SIZE, - skiprows=self.skip_rows, - ): - break + chunk = next( + pandas.read_csv( + path_to_file, + chunksize=self.cfg.DATA.LOADER_CHUNK_SIZE, + skiprows=self.skip_rows, + ) + ) except Exception: self.skip_rows = 0 return self._get_chunk(path_to_file, chunksize) @@ -193,7 +192,8 @@ def __getitem__(self, index): index, self._num_yielded = index if self.cfg.MULTIGRID.SHORT_CYCLE: index, short_cycle_idx = index - + if self.dummy_output is not None: + return self.dummy_output if self.mode in ["train", "val"]: # -1 indicates random sampling. temporal_sample_index = -1 @@ -212,25 +212,17 @@ def __getitem__(self, index): # Decreasing the scale is equivalent to using a larger "span" # in a sampling grid. min_scale = int( - round( - float(min_scale) - * crop_size - / self.cfg.MULTIGRID.DEFAULT_S - ) + round(float(min_scale) * crop_size / self.cfg.MULTIGRID.DEFAULT_S) ) elif self.mode in ["test"]: temporal_sample_index = ( - self._spatial_temporal_idx[index] - // self.cfg.TEST.NUM_SPATIAL_CROPS + self._spatial_temporal_idx[index] // self.cfg.TEST.NUM_SPATIAL_CROPS ) # spatial_sample_index is in [0, 1, 2]. Corresponding to left, # center, or right if width is larger than height, and top, middle, # or bottom if height is larger than width. spatial_sample_index = ( - ( - self._spatial_temporal_idx[index] - % self.cfg.TEST.NUM_SPATIAL_CROPS - ) + (self._spatial_temporal_idx[index] % self.cfg.TEST.NUM_SPATIAL_CROPS) if self.cfg.TEST.NUM_SPATIAL_CROPS > 1 else 1 ) @@ -244,13 +236,9 @@ def __getitem__(self, index): # min_scale, max_scale, and crop_size are expect to be the same. assert len({min_scale, max_scale}) == 1 else: - raise NotImplementedError( - "Does not support {} mode".format(self.mode) - ) + raise NotImplementedError("Does not support {} mode".format(self.mode)) num_decode = ( - self.cfg.DATA.TRAIN_CROP_NUM_TEMPORAL - if self.mode in ["train"] - else 1 + self.cfg.DATA.TRAIN_CROP_NUM_TEMPORAL if self.mode in ["train"] else 1 ) min_scale, max_scale, crop_size = [min_scale], [max_scale], [crop_size] if len(min_scale) < num_decode: @@ -262,10 +250,8 @@ def __getitem__(self, index): ) crop_size += ( [self.cfg.MULTIGRID.DEFAULT_S] * (num_decode - len(crop_size)) - if self.cfg.MULTIGRID.LONG_CYCLE - or self.cfg.MULTIGRID.SHORT_CYCLE - else [self.cfg.DATA.TRAIN_CROP_SIZE] - * (num_decode - len(crop_size)) + if self.cfg.MULTIGRID.LONG_CYCLE or self.cfg.MULTIGRID.SHORT_CYCLE + else [self.cfg.DATA.TRAIN_CROP_SIZE] * (num_decode - len(crop_size)) ) assert self.mode in ["train", "val"] # Try to decode and sample a clip from a video. If the video can not be @@ -284,7 +270,10 @@ def __getitem__(self, index): self._path_to_videos[index], e ) ) - # Select a random video if the current video was not able to access. + if self.mode not in ["test"]: + # let's try another one + index = random.randint(0, len(self._path_to_videos) - 1) + continue # Select a random video if the current video was not able to access. if video_container is None: logger.warning( "Failed to meta load video idx {} from {}; trial {}".format( @@ -310,35 +299,22 @@ def __getitem__(self, index): sampling_rate = [sampling_rate] if len(num_frames) < num_decode: num_frames.extend( - [ - num_frames[-1] - for i in range(num_decode - len(num_frames)) - ] + [num_frames[-1] for i in range(num_decode - len(num_frames))] ) # base case where keys have same frame-rate as query sampling_rate.extend( - [ - sampling_rate[-1] - for i in range(num_decode - len(sampling_rate)) - ] + [sampling_rate[-1] for i in range(num_decode - len(sampling_rate))] ) elif len(num_frames) > num_decode: num_frames = num_frames[:num_decode] sampling_rate = sampling_rate[:num_decode] if self.mode in ["train"]: - assert ( - len(min_scale) - == len(max_scale) - == len(crop_size) - == num_decode - ) + assert len(min_scale) == len(max_scale) == len(crop_size) == num_decode target_fps = self.cfg.DATA.TARGET_FPS if self.cfg.DATA.TRAIN_JITTER_FPS > 0.0 and self.mode in ["train"]: - target_fps += random.uniform( - 0.0, self.cfg.DATA.TRAIN_JITTER_FPS - ) + target_fps += random.uniform(0.0, self.cfg.DATA.TRAIN_JITTER_FPS) # Decode video. Meta info is used to perform selective decoding. frames, time_idx, tdiff = decoder.decode( @@ -347,18 +323,16 @@ def __getitem__(self, index): num_frames, temporal_sample_index, self.cfg.TEST.NUM_ENSEMBLE_VIEWS, - video_meta=self._video_meta[index] - if len(self._video_meta) < 5e6 - else {}, # do not cache on huge datasets + video_meta=( + self._video_meta[index] if len(self._video_meta) < 5e6 else {} + ), # do not cache on huge datasets target_fps=target_fps, backend=self.cfg.DATA.DECODING_BACKEND, use_offset=self.cfg.DATA.USE_OFFSET_SAMPLING, - max_spatial_scale=min_scale[0] - if all(x == min_scale[0] for x in min_scale) - else 0, # if self.mode in ["test"] else 0, - time_diff_prob=self.p_convert_dt - if self.mode in ["train"] - else 0.0, + max_spatial_scale=( + min_scale[0] if all(x == min_scale[0] for x in min_scale) else 0 + ), # if self.mode in ["test"] else 0, + time_diff_prob=self.p_convert_dt if self.mode in ["train"] else 0.0, temporally_rnd_clips=True, min_delta=self.cfg.CONTRASTIVE.DELTA_CLIPS_MIN, max_delta=self.cfg.CONTRASTIVE.DELTA_CLIPS_MAX, @@ -401,10 +375,7 @@ def __getitem__(self, index): f_out[idx] = f_out[idx].float() f_out[idx] = f_out[idx] / 255.0 - if ( - self.mode in ["train"] - and self.cfg.DATA.SSL_COLOR_JITTER - ): + if self.mode in ["train"] and self.cfg.DATA.SSL_COLOR_JITTER: f_out[idx] = transform.color_jitter_video_ssl( f_out[idx], bri_con_sat=self.cfg.DATA.SSL_COLOR_BRI_CON_SAT, @@ -415,7 +386,7 @@ def __getitem__(self, index): gaussan_sigma_max=self.cfg.DATA.SSL_BLUR_SIGMA_MAX, ) - if self.randaug: + if self.aug and self.cfg.AUG.AA_TYPE: aug_transform = create_random_augment( input_size=(f_out[idx].size(1), f_out[idx].size(2)), auto_augment=self.cfg.AUG.AA_TYPE, @@ -441,14 +412,10 @@ def __getitem__(self, index): self.cfg.DATA.TRAIN_JITTER_ASPECT_RELATIVE, ) relative_scales = ( - None - if (self.mode not in ["train"] or len(scl) == 0) - else scl + None if (self.mode not in ["train"] or len(scl) == 0) else scl ) relative_aspect = ( - None - if (self.mode not in ["train"] or len(asp) == 0) - else asp + None if (self.mode not in ["train"] or len(asp) == 0) else asp ) f_out[idx] = utils.spatial_sampling( f_out[idx], @@ -460,9 +427,11 @@ def __getitem__(self, index): inverse_uniform_sampling=self.cfg.DATA.INV_UNIFORM_SAMPLE, aspect_ratio=relative_aspect, scale=relative_scales, - motion_shift=self.cfg.DATA.TRAIN_JITTER_MOTION_SHIFT - if self.mode in ["train"] - else False, + motion_shift=( + self.cfg.DATA.TRAIN_JITTER_MOTION_SHIFT + if self.mode in ["train"] + else False + ), ) if self.rand_erase: @@ -478,23 +447,64 @@ def __getitem__(self, index): ).permute(1, 0, 2, 3) f_out[idx] = utils.pack_pathway_output(self.cfg, f_out[idx]) + if self.cfg.AUG.GEN_MASK_LOADER: + mask = self._gen_mask() + f_out[idx] = f_out[idx] + [torch.Tensor(), mask] frames = f_out[0] if num_out == 1 else f_out time_idx = np.array(time_idx_out) - if num_aug > 1: - label = [label] * num_aug - index = [index] * num_aug + if ( + num_aug * num_decode > 1 + and not self.cfg.MODEL.MODEL_NAME == "ContrastiveModel" + ): + label = [label] * num_aug * num_decode + index = [index] * num_aug * num_decode + if self.cfg.DATA.DUMMY_LOAD: + if self.dummy_output is None: + self.dummy_output = (frames, label, index, time_idx, {}) return frames, label, index, time_idx, {} else: - raise RuntimeError( - "Failed to fetch video idx {} from {}; after {} trials".format( - index, self._path_to_videos[index], i_try - ) + logger.warning( + "Failed to fetch video after {} retries.".format(self._num_retries) + ) + + def _gen_mask(self): + if self.cfg.AUG.MASK_TUBE: + num_masking_patches = round( + np.prod(self.cfg.AUG.MASK_WINDOW_SIZE) * self.cfg.AUG.MASK_RATIO + ) + min_mask = num_masking_patches // 5 + masked_position_generator = MaskingGenerator( + mask_window_size=self.cfg.AUG.MASK_WINDOW_SIZE, + num_masking_patches=num_masking_patches, + max_num_patches=None, + min_num_patches=min_mask, + ) + mask = masked_position_generator() + mask = np.tile(mask, (8, 1, 1)) + elif self.cfg.AUG.MASK_FRAMES: + mask = np.zeros(shape=self.cfg.AUG.MASK_WINDOW_SIZE, dtype=int) + n_mask = round(self.cfg.AUG.MASK_WINDOW_SIZE[0] * self.cfg.AUG.MASK_RATIO) + mask_t_ind = random.sample( + range(0, self.cfg.AUG.MASK_WINDOW_SIZE[0]), n_mask + ) + mask[mask_t_ind, :, :] += 1 + else: + num_masking_patches = round( + np.prod(self.cfg.AUG.MASK_WINDOW_SIZE) * self.cfg.AUG.MASK_RATIO + ) + max_mask = np.prod(self.cfg.AUG.MASK_WINDOW_SIZE[1:]) + min_mask = max_mask // 5 + masked_position_generator = MaskingGenerator3D( + mask_window_size=self.cfg.AUG.MASK_WINDOW_SIZE, + num_masking_patches=num_masking_patches, + max_num_patches=max_mask, + min_num_patches=min_mask, ) + mask = masked_position_generator() + return mask def _frame_to_list_img(self, frames): - img_list = [ - transforms.ToPILImage()(frames[i]) for i in range(frames.size(0)) - ] + img_list = [transforms.ToPILImage()(frames[i]) for i in range(frames.size(0))] return img_list def _list_img_to_frames(self, img_list): diff --git a/slowfast/datasets/loader.py b/slowfast/datasets/loader.py index 71c4aaf66..712913c8a 100644 --- a/slowfast/datasets/loader.py +++ b/slowfast/datasets/loader.py @@ -4,15 +4,14 @@ """Data loader.""" import itertools -import numpy as np from functools import partial -from typing import List + +import numpy as np import torch +from slowfast.datasets.multigrid_helper import ShortCycleBatchSampler from torch.utils.data._utils.collate import default_collate from torch.utils.data.distributed import DistributedSampler -from torch.utils.data.sampler import RandomSampler, Sampler - -from slowfast.datasets.multigrid_helper import ShortCycleBatchSampler +from torch.utils.data.sampler import RandomSampler from . import utils as utils from .build import build_dataset @@ -75,9 +74,9 @@ def detection_collate(batch): bboxes = np.concatenate(bboxes, axis=0) collated_extra_data[key] = torch.tensor(bboxes).float() elif key == "metadata": - collated_extra_data[key] = torch.tensor( - list(itertools.chain(*data)) - ).view(-1, 2) + collated_extra_data[key] = torch.tensor(list(itertools.chain(*data))).view( + -1, 2 + ) else: collated_extra_data[key] = default_collate(data) @@ -124,11 +123,7 @@ def construct_loader(cfg, split, is_precise_bn=False): worker_init_fn=utils.loader_worker_init_fn(dataset), ) else: - if ( - cfg.MULTIGRID.SHORT_CYCLE - and split in ["train"] - and not is_precise_bn - ): + if cfg.MULTIGRID.SHORT_CYCLE and split in ["train"] and not is_precise_bn: # Create a sampler for multi-process training sampler = utils.create_sampler(dataset, shuffle, cfg) batch_sampler = ShortCycleBatchSampler( @@ -148,13 +143,20 @@ def construct_loader(cfg, split, is_precise_bn=False): # Create a loader if cfg.DETECTION.ENABLE: collate_func = detection_collate - elif cfg.AUG.NUM_SAMPLE > 1 and split in ["train"]: + elif ( + ( + cfg.AUG.NUM_SAMPLE > 1 + or cfg.DATA.TRAIN_CROP_NUM_TEMPORAL > 1 + or cfg.DATA.TRAIN_CROP_NUM_SPATIAL > 1 + ) + and split in ["train"] + and not cfg.MODEL.MODEL_NAME == "ContrastiveModel" + ): collate_func = partial( multiple_samples_collate, fold="imagenet" in dataset_name ) else: collate_func = None - loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, @@ -176,10 +178,7 @@ def shuffle_dataset(loader, cur_epoch): loader (loader): data loader to perform shuffle. cur_epoch (int): number of the current epoch. """ - if ( - loader._dataset_kind - == torch.utils.data.dataloader._DatasetKind.Iterable - ): + if loader._dataset_kind == torch.utils.data.dataloader._DatasetKind.Iterable: if hasattr(loader.dataset, "sampler"): sampler = loader.dataset.sampler else: @@ -192,10 +191,17 @@ def shuffle_dataset(loader, cur_epoch): if isinstance(loader.batch_sampler, ShortCycleBatchSampler) else loader.sampler ) - assert isinstance( - sampler, (RandomSampler, DistributedSampler) - ), "Sampler type '{}' not supported".format(type(sampler)) + assert isinstance(sampler, (RandomSampler, DistributedSampler)), ( + "Sampler type '{}' not supported".format(type(sampler)) + ) # RandomSampler handles shuffling automatically if isinstance(sampler, DistributedSampler): # DistributedSampler shuffles data based on epoch sampler.set_epoch(cur_epoch) + + if hasattr(loader.dataset, "prefetcher"): + sampler = loader.dataset.prefetcher.sampler + if isinstance(sampler, DistributedSampler): + # DistributedSampler shuffles data based on epoch + print("prefetcher sampler") + sampler.set_epoch(cur_epoch) diff --git a/slowfast/datasets/mixup.py b/slowfast/datasets/mixup.py index a75846d04..f04af6795 100644 --- a/slowfast/datasets/mixup.py +++ b/slowfast/datasets/mixup.py @@ -181,9 +181,12 @@ def _mix_batch(self, x): return lam def __call__(self, x, target): - assert len(x) > 1, "Batch size should be greater than 1 for mixup." - lam = self._mix_batch(x) - target = mixup_target( - target, self.num_classes, lam, self.label_smoothing - ) + if self.mix_prob > 0.0: + assert len(x) > 1, "Batch size should be greater than 1 for mixup." + lam = self._mix_batch(x) + elif self.mix_prob == 0.0: + lam = 1.0 + else: + raise NotImplementedError + target = mixup_target(target, self.num_classes, lam, self.label_smoothing) return x, target diff --git a/slowfast/datasets/multigrid_helper.py b/slowfast/datasets/multigrid_helper.py index 2381a4c9b..afce4ecbf 100644 --- a/slowfast/datasets/multigrid_helper.py +++ b/slowfast/datasets/multigrid_helper.py @@ -40,8 +40,9 @@ def __init__(self, sampler, batch_size, drop_last, cfg): ) if not isinstance(drop_last, bool): raise ValueError( - "drop_last should be a boolean value, but got " - "drop_last={}".format(drop_last) + "drop_last should be a boolean value, but got drop_last={}".format( + drop_last + ) ) self.sampler = sampler self.drop_last = drop_last @@ -49,10 +50,7 @@ def __init__(self, sampler, batch_size, drop_last, cfg): bs_factor = [ int( round( - ( - float(cfg.DATA.TRAIN_CROP_SIZE) - / (s * cfg.MULTIGRID.DEFAULT_S) - ) + (float(cfg.DATA.TRAIN_CROP_SIZE) / (s * cfg.MULTIGRID.DEFAULT_S)) ** 2 ) ) diff --git a/slowfast/datasets/ptv_datasets.py b/slowfast/datasets/ptv_datasets.py index e5de94bc8..a3b5fd710 100644 --- a/slowfast/datasets/ptv_datasets.py +++ b/slowfast/datasets/ptv_datasets.py @@ -3,27 +3,10 @@ import functools import os from typing import Dict -import torch -from torch.utils.data import ( - DistributedSampler, - RandomSampler, - SequentialSampler, -) -from torchvision.transforms import Compose, Lambda -from torchvision.transforms._transforms_video import ( - NormalizeVideo, - RandomCropVideo, - RandomHorizontalFlipVideo, -) import slowfast.utils.logging as logging - -from pytorchvideo.data import ( - Charades, - LabeledVideoDataset, - SSv2, - make_clip_sampler, -) +import torch +from pytorchvideo.data import Charades, LabeledVideoDataset, make_clip_sampler, SSv2 from pytorchvideo.data.labeled_video_paths import LabeledVideoPaths from pytorchvideo.transforms import ( ApplyTransformToKey, @@ -32,6 +15,13 @@ UniformCropVideo, UniformTemporalSubsample, ) +from torch.utils.data import DistributedSampler, RandomSampler, SequentialSampler +from torchvision.transforms import Compose, Lambda +from torchvision.transforms._transforms_video import ( + NormalizeVideo, + RandomCropVideo, + RandomHorizontalFlipVideo, +) from . import utils as utils from .build import DATASET_REGISTRY @@ -170,12 +160,8 @@ def Ptvkinetics(cfg, mode): logger.info("Constructing Ptvkinetics {}...".format(mode)) - clip_duration = ( - cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE / cfg.DATA.TARGET_FPS - ) - path_to_file = os.path.join( - cfg.DATA.PATH_TO_DATA_DIR, "{}.csv".format(mode) - ) + clip_duration = cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE / cfg.DATA.TARGET_FPS + path_to_file = os.path.join(cfg.DATA.PATH_TO_DATA_DIR, "{}.csv".format(mode)) labeled_video_paths = LabeledVideoPaths.from_path(path_to_file) num_videos = len(labeled_video_paths) labeled_video_paths.path_prefix = cfg.DATA.PATH_PREFIX @@ -220,9 +206,7 @@ def Ptvkinetics(cfg, mode): if cfg.NUM_GPUS > 1: video_sampler = DistributedSampler else: - video_sampler = ( - RandomSampler if mode == "train" else SequentialSampler - ) + video_sampler = RandomSampler if mode == "train" else SequentialSampler else: num_clips = cfg.TEST.NUM_ENSEMBLE_VIEWS num_crops = cfg.TEST.NUM_SPATIAL_CROPS @@ -236,9 +220,7 @@ def Ptvkinetics(cfg, mode): UniformTemporalSubsample(cfg.DATA.NUM_FRAMES), Lambda(div255), NormalizeVideo(cfg.DATA.MEAN, cfg.DATA.STD), - ShortSideScale( - size=cfg.DATA.TRAIN_JITTER_SCALES[0] - ), + ShortSideScale(size=cfg.DATA.TRAIN_JITTER_SCALES[0]), ] ), ), @@ -253,9 +235,7 @@ def Ptvkinetics(cfg, mode): num_clips, num_crops, ) - video_sampler = ( - DistributedSampler if cfg.NUM_GPUS > 1 else SequentialSampler - ) + video_sampler = DistributedSampler if cfg.NUM_GPUS > 1 else SequentialSampler return PTVDatasetWrapper( num_videos=num_videos, @@ -285,11 +265,7 @@ def process_charades_label(x, mode, num_classes): Returns: x (dict): video clip with updated label information. """ - label = ( - utils.aggregate_labels(x["label"]) - if mode == "train" - else x["video_label"] - ) + label = utils.aggregate_labels(x["label"]) if mode == "train" else x["video_label"] x["label"] = torch.as_tensor(utils.as_binary_vector(label, num_classes)) return x @@ -381,9 +357,7 @@ def Ptvcharades(cfg, mode): if cfg.NUM_GPUS > 1: video_sampler = DistributedSampler else: - video_sampler = ( - RandomSampler if mode == "train" else SequentialSampler - ) + video_sampler = RandomSampler if mode == "train" else SequentialSampler else: num_clips = cfg.TEST.NUM_ENSEMBLE_VIEWS num_crops = cfg.TEST.NUM_SPATIAL_CROPS @@ -423,9 +397,7 @@ def Ptvcharades(cfg, mode): num_clips, num_crops, ) - video_sampler = ( - DistributedSampler if cfg.NUM_GPUS > 1 else SequentialSampler - ) + video_sampler = DistributedSampler if cfg.NUM_GPUS > 1 else SequentialSampler data_path = os.path.join(cfg.DATA.PATH_TO_DATA_DIR, "{}.csv".format(mode)) dataset = Charades( @@ -516,9 +488,7 @@ def Ptvssv2(cfg, mode): if cfg.NUM_GPUS > 1: video_sampler = DistributedSampler else: - video_sampler = ( - RandomSampler if mode == "train" else SequentialSampler - ) + video_sampler = RandomSampler if mode == "train" else SequentialSampler else: assert cfg.TEST.NUM_ENSEMBLE_VIEWS == 1 num_clips = cfg.TEST.NUM_ENSEMBLE_VIEWS @@ -552,9 +522,7 @@ def Ptvssv2(cfg, mode): num_clips, num_crops, ) - video_sampler = ( - DistributedSampler if cfg.NUM_GPUS > 1 else SequentialSampler - ) + video_sampler = DistributedSampler if cfg.NUM_GPUS > 1 else SequentialSampler label_name_file = os.path.join( cfg.DATA.PATH_TO_DATA_DIR, "something-something-v2-labels.json" diff --git a/slowfast/datasets/rand_augment.py b/slowfast/datasets/rand_augment.py index 8ee81fc94..c3f741a9b 100644 --- a/slowfast/datasets/rand_augment.py +++ b/slowfast/datasets/rand_augment.py @@ -27,9 +27,10 @@ """ import math -import numpy as np import random import re + +import numpy as np import PIL from PIL import Image, ImageEnhance, ImageOps @@ -65,46 +66,34 @@ def _check_args_tf(kwargs): def shear_x(img, factor, **kwargs): _check_args_tf(kwargs) - return img.transform( - img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs - ) + return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs) def shear_y(img, factor, **kwargs): _check_args_tf(kwargs) - return img.transform( - img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs - ) + return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs) def translate_x_rel(img, pct, **kwargs): pixels = pct * img.size[0] _check_args_tf(kwargs) - return img.transform( - img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs - ) + return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs) def translate_y_rel(img, pct, **kwargs): pixels = pct * img.size[1] _check_args_tf(kwargs) - return img.transform( - img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs - ) + return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs) def translate_x_abs(img, pixels, **kwargs): _check_args_tf(kwargs) - return img.transform( - img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs - ) + return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs) def translate_y_abs(img, pixels, **kwargs): _check_args_tf(kwargs) - return img.transform( - img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs - ) + return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs) def rotate(img, degrees, **kwargs): @@ -349,12 +338,12 @@ def __init__(self, name, prob=0.5, magnitude=10, hparams=None): self.magnitude = magnitude self.hparams = hparams.copy() self.kwargs = { - "fillcolor": hparams["img_mean"] - if "img_mean" in hparams - else _FILL, - "resample": hparams["interpolation"] - if "interpolation" in hparams - else _RANDOM_INTERPOLATION, + "fillcolor": hparams["img_mean"] if "img_mean" in hparams else _FILL, + "resample": ( + hparams["interpolation"] + if "interpolation" in hparams + else _RANDOM_INTERPOLATION + ), } # If magnitude_std is > 0, we introduce some randomness @@ -371,15 +360,11 @@ def __call__(self, img_list): magnitude = random.gauss(magnitude, self.magnitude_std) magnitude = min(_MAX_LEVEL, max(0, magnitude)) # clip to valid range level_args = ( - self.level_fn(magnitude, self.hparams) - if self.level_fn is not None - else () + self.level_fn(magnitude, self.hparams) if self.level_fn is not None else () ) if isinstance(img_list, list): - return [ - self.aug_fn(img, *level_args, **self.kwargs) for img in img_list - ] + return [self.aug_fn(img, *level_args, **self.kwargs) for img in img_list] else: return self.aug_fn(img_list, *level_args, **self.kwargs) @@ -527,7 +512,5 @@ def rand_augment_transform(config_str, hparams): ra_ops = rand_augment_ops( magnitude=magnitude, hparams=hparams, transforms=transforms ) - choice_weights = ( - None if weight_idx is None else _select_rand_weights(weight_idx) - ) + choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx) return RandAugment(ra_ops, num_layers, choice_weights=choice_weights) diff --git a/slowfast/datasets/random_erasing.py b/slowfast/datasets/random_erasing.py index 952060f94..755790b75 100644 --- a/slowfast/datasets/random_erasing.py +++ b/slowfast/datasets/random_erasing.py @@ -10,23 +10,21 @@ Copyright Zhun Zhong & Liang Zheng Hacked together by / Copyright 2020 Ross Wightman """ + import math import random + import torch -def _get_pixels( - per_pixel, rand_color, patch_size, dtype=torch.float32, device="cuda" -): +def _get_pixels(per_pixel, rand_color, patch_size, dtype=torch.float32, device="cuda"): # NOTE I've seen CUDA illegal memory access errors being caused by the normal_() # paths, flip the order so normal is run on CPU if this becomes a problem # Issue has been fixed in master https://github.com/pytorch/pytorch/issues/19508 if per_pixel: return torch.empty(patch_size, dtype=dtype, device=device).normal_() elif rand_color: - return torch.empty( - (patch_size[0], 1, 1), dtype=dtype, device=device - ).normal_() + return torch.empty((patch_size[0], 1, 1), dtype=dtype, device=device).normal_() else: return torch.zeros((patch_size[0], 1, 1), dtype=dtype, device=device) @@ -144,9 +142,7 @@ def _erase_cube( left = random.randint(0, img_w - w) for i in range(batch_start, batch_size): img_instance = img[i] - img_instance[ - :, top : top + h, left : left + w - ] = _get_pixels( + img_instance[:, top : top + h, left : left + w] = _get_pixels( self.per_pixel, self.rand_color, (chan, h, w), @@ -161,9 +157,7 @@ def __call__(self, input): else: batch_size, chan, img_h, img_w = input.size() # skip first slice of batch if num_splits is set (for clean portion of samples) - batch_start = ( - batch_size // self.num_splits if self.num_splits > 1 else 0 - ) + batch_start = batch_size // self.num_splits if self.num_splits > 1 else 0 if self.cube: self._erase_cube( input, diff --git a/slowfast/datasets/ssv2.py b/slowfast/datasets/ssv2.py index e5715d69d..dc2314882 100644 --- a/slowfast/datasets/ssv2.py +++ b/slowfast/datasets/ssv2.py @@ -2,14 +2,14 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import json -import numpy as np import os import random from itertools import chain as chain -import torch -import torch.utils.data +import numpy as np import slowfast.utils.logging as logging +import torch +import torch.utils.data from slowfast.utils.env import pathmgr from . import utils as utils @@ -63,13 +63,21 @@ def __init__(self, cfg, mode, num_retries=10): if self.mode in ["train", "val"]: self._num_clips = 1 elif self.mode in ["test"]: - self._num_clips = ( - cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS - ) + self._num_clips = cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS logger.info("Constructing Something-Something V2 {}...".format(mode)) self._construct_loader() + self.aug = False + self.rand_erase = False + self.use_temporal_gradient = False + self.temporal_gradient_rate = 0.0 + + if self.mode == "train" and self.cfg.AUG.ENABLE: + self.aug = True + if self.cfg.AUG.RE_PROB > 0: + self.rand_erase = True + def _construct_loader(self): """ Construct the video loader. @@ -109,9 +117,7 @@ def _construct_loader(self): self.cfg.DATA.PATH_TO_DATA_DIR, "{}.csv".format("train" if self.mode == "train" else "val"), ) - assert pathmgr.exists(path_to_file), "{} dir not found".format( - path_to_file - ) + assert pathmgr.exists(path_to_file), "{} dir not found".format(path_to_file) self._path_to_videos, _ = utils.load_image_lists( path_to_file, self.cfg.DATA.PATH_PREFIX @@ -134,24 +140,18 @@ def _construct_loader(self): # Extend self when self._num_clips > 1 (during testing). self._path_to_videos = list( - chain.from_iterable( - [[x] * self._num_clips for x in self._path_to_videos] - ) + chain.from_iterable([[x] * self._num_clips for x in self._path_to_videos]) ) self._labels = list( chain.from_iterable([[x] * self._num_clips for x in self._labels]) ) self._spatial_temporal_idx = list( chain.from_iterable( - [ - range(self._num_clips) - for _ in range(len(self._path_to_videos)) - ] + [range(self._num_clips) for _ in range(len(self._path_to_videos))] ) ) logger.info( - "Something-Something V2 dataloader constructed " - " (size: {}) from {}".format( + "Something-Something V2 dataloader constructed (size: {}) from {}".format( len(self._path_to_videos), path_to_file ) ) @@ -215,28 +215,21 @@ def __getitem__(self, index): # Decreasing the scale is equivalent to using a larger "span" # in a sampling grid. min_scale = int( - round( - float(min_scale) - * crop_size - / self.cfg.MULTIGRID.DEFAULT_S - ) + round(float(min_scale) * crop_size / self.cfg.MULTIGRID.DEFAULT_S) ) elif self.mode in ["test"]: # spatial_sample_index is in [0, 1, 2]. Corresponding to left, # center, or right if width is larger than height, and top, middle, # or bottom if height is larger than width. spatial_sample_index = ( - self._spatial_temporal_idx[index] - % self.cfg.TEST.NUM_SPATIAL_CROPS + self._spatial_temporal_idx[index] % self.cfg.TEST.NUM_SPATIAL_CROPS ) min_scale, max_scale, crop_size = [self.cfg.DATA.TEST_CROP_SIZE] * 3 # The testing is deterministic and no jitter should be performed. # min_scale, max_scale, and crop_size are expect to be the same. assert len({min_scale, max_scale, crop_size}) == 1 else: - raise NotImplementedError( - "Does not support {} mode".format(self.mode) - ) + raise NotImplementedError("Does not support {} mode".format(self.mode)) label = self._labels[index] @@ -249,23 +242,63 @@ def __getitem__(self, index): ) ) - # Perform color normalization. - frames = utils.tensor_normalize( - frames, self.cfg.DATA.MEAN, self.cfg.DATA.STD - ) + if self.aug: + if self.cfg.AUG.NUM_SAMPLE > 1: + frame_list = [] + label_list = [] + index_list = [] + for _ in range(self.cfg.AUG.NUM_SAMPLE): + new_frames = utils.aug_frame( + self.cfg, + self.mode, + self.rand_erase, + frames, + spatial_sample_index, + min_scale, + max_scale, + crop_size, + ) + new_frames = utils.pack_pathway_output(self.cfg, new_frames) + frame_list.append(new_frames) + label_list.append(label) + index_list.append(index) + return ( + frame_list, + label_list, + index_list, + [0] * self.cfg.AUG.NUM_SAMPLE, + {}, + ) - # T H W C -> C T H W. - frames = frames.permute(3, 0, 1, 2) - # Perform data augmentation. - frames = utils.spatial_sampling( - frames, - spatial_idx=spatial_sample_index, - min_scale=min_scale, - max_scale=max_scale, - crop_size=crop_size, - random_horizontal_flip=self.cfg.DATA.RANDOM_FLIP, - inverse_uniform_sampling=self.cfg.DATA.INV_UNIFORM_SAMPLE, - ) + else: + frames = utils.aug_frame( + self.cfg, + self.mode, + self.rand_erase, + frames, + spatial_sample_index, + min_scale, + max_scale, + crop_size, + ) + else: + # Perform color normalization. + frames = utils.tensor_normalize( + frames, self.cfg.DATA.MEAN, self.cfg.DATA.STD + ) + + # T H W C -> C T H W. + frames = frames.permute(3, 0, 1, 2) + # Perform data augmentation. + frames = utils.spatial_sampling( + frames, + spatial_idx=spatial_sample_index, + min_scale=min_scale, + max_scale=max_scale, + crop_size=crop_size, + random_horizontal_flip=self.cfg.DATA.RANDOM_FLIP, + inverse_uniform_sampling=self.cfg.DATA.INV_UNIFORM_SAMPLE, + ) frames = utils.pack_pathway_output(self.cfg, frames) return frames, label, index, 0, {} diff --git a/slowfast/datasets/transform.py b/slowfast/datasets/transform.py index a95b21a2e..f81e449d4 100644 --- a/slowfast/datasets/transform.py +++ b/slowfast/datasets/transform.py @@ -3,10 +3,11 @@ import logging import math -import numpy as np # import cv2 import random + +import numpy as np import torch import torchvision as tv import torchvision.transforms.functional as F @@ -67,17 +68,13 @@ def random_short_side_scale_jitter( `num boxes` x 4. """ if inverse_uniform_sampling: - size = int( - round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size)) - ) + size = int(round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size))) else: size = int(round(np.random.uniform(min_size, max_size))) height = images.shape[2] width = images.shape[3] - if (width <= height and width == size) or ( - height <= width and height == size - ): + if (width <= height and width == size) or (height <= width and height == size): return images, boxes new_width = size new_height = size @@ -145,13 +142,9 @@ def random_crop(images, size, boxes=None): x_offset = 0 if width > size: x_offset = int(np.random.randint(0, width - size)) - cropped = images[ - :, :, y_offset : y_offset + size, x_offset : x_offset + size - ] + cropped = images[:, :, y_offset : y_offset + size, x_offset : x_offset + size] - cropped_boxes = ( - crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None - ) + cropped_boxes = crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None return cropped, cropped_boxes @@ -243,12 +236,8 @@ def uniform_crop(images, size, spatial_idx, boxes=None, scale_size=None): x_offset = 0 elif spatial_idx == 2: x_offset = width - size - cropped = images[ - :, :, y_offset : y_offset + size, x_offset : x_offset + size - ] - cropped_boxes = ( - crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None - ) + cropped = images[:, :, y_offset : y_offset + size, x_offset : x_offset + size] + cropped_boxes = crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None if ndim == 3: cropped = cropped.squeeze(0) return cropped, cropped_boxes @@ -305,9 +294,7 @@ def grayscale(images): """ # R -> 0.299, G -> 0.587, B -> 0.114. img_gray = torch.tensor(images) - gray_channel = ( - 0.299 * images[:, 2] + 0.587 * images[:, 1] + 0.114 * images[:, 0] - ) + gray_channel = 0.299 * images[:, 2] + 0.587 * images[:, 1] + 0.114 * images[:, 0] img_gray[:, 0] = gray_channel img_gray[:, 1] = gray_channel img_gray[:, 2] = gray_channel @@ -448,9 +435,7 @@ def lighting_jitter(images, alphastd, eigval, eigvec): elif len(images.shape) == 4: out_images[:, idx] = images[:, idx] + rgb[2 - idx] else: - raise NotImplementedError( - f"Unsupported dimension {len(images.shape)}" - ) + raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") return out_images @@ -469,19 +454,11 @@ def color_normalization(images, mean, stddev): `num frames` x `channel` x `height` x `width`. """ if len(images.shape) == 3: - assert ( - len(mean) == images.shape[0] - ), "channel mean not computed properly" - assert ( - len(stddev) == images.shape[0] - ), "channel stddev not computed properly" + assert len(mean) == images.shape[0], "channel mean not computed properly" + assert len(stddev) == images.shape[0], "channel stddev not computed properly" elif len(images.shape) == 4: - assert ( - len(mean) == images.shape[1] - ), "channel mean not computed properly" - assert ( - len(stddev) == images.shape[1] - ), "channel stddev not computed properly" + assert len(mean) == images.shape[1], "channel mean not computed properly" + assert len(stddev) == images.shape[1], "channel stddev not computed properly" else: raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") @@ -493,9 +470,7 @@ def color_normalization(images, mean, stddev): elif len(images.shape) == 4: out_images[:, idx] = (images[:, idx] - mean[idx]) / stddev[idx] else: - raise NotImplementedError( - f"Unsupported dimension {len(images.shape)}" - ) + raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") return out_images @@ -651,9 +626,7 @@ def create_random_augment( if interpolation and interpolation != "random": aa_params["interpolation"] = _pil_interp(interpolation) if auto_augment.startswith("rand"): - return transforms.Compose( - [rand_augment_transform(auto_augment, aa_params)] - ) + return transforms.Compose([rand_augment_transform(auto_augment, aa_params)]) raise NotImplementedError @@ -667,9 +640,7 @@ def random_sized_crop_img( """ Performs Inception-style cropping (used for training). """ - assert ( - len(im.shape) == 3 - ), "Currently only support image for random_sized_crop" + assert len(im.shape) == 3, "Currently only support image for random_sized_crop" h, w = im.shape[1:3] i, j, h, w = _get_param_spatial_crop( scale=jitter_scale, @@ -789,16 +760,191 @@ def __repr__(self): else: interpolate_str = _pil_interpolation_to_str[self.interpolation] format_string = self.__class__.__name__ + "(size={0}".format(self.size) - format_string += ", scale={0}".format( - tuple(round(s, 4) for s in self.scale) - ) - format_string += ", ratio={0}".format( - tuple(round(r, 4) for r in self.ratio) - ) + format_string += ", scale={0}".format(tuple(round(s, 4) for s in self.scale)) + format_string += ", ratio={0}".format(tuple(round(r, 4) for r in self.ratio)) format_string += ", interpolation={0})".format(interpolate_str) return format_string +""" +This implementation is based on +https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py +Licensed under The MIT License +""" + + +class MaskingGenerator: + def __init__( + self, + mask_window_size, + num_masking_patches, + min_num_patches=16, + max_num_patches=None, + min_aspect=0.3, + max_aspect=None, + ): + if not isinstance( + mask_window_size, + ( + list, + tuple, + ), + ): + mask_window_size = (mask_window_size,) * 2 + self.height, self.width = mask_window_size + + self.num_patches = self.height * self.width + self.num_masking_patches = num_masking_patches + + self.min_num_patches = min_num_patches + self.max_num_patches = ( + num_masking_patches if max_num_patches is None else max_num_patches + ) + + max_aspect = max_aspect or 1 / min_aspect + self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) + + def __repr__(self): + repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % ( + self.height, + self.width, + self.min_num_patches, + self.max_num_patches, + self.num_masking_patches, + self.log_aspect_ratio[0], + self.log_aspect_ratio[1], + ) + return repr_str + + def get_shape(self): + return self.height, self.width + + def _mask(self, mask, max_mask_patches): + delta = 0 + for _ in range(10): + target_area = random.uniform(self.min_num_patches, max_mask_patches) + aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) + h = int(round(math.sqrt(target_area * aspect_ratio))) + w = int(round(math.sqrt(target_area / aspect_ratio))) + if w < self.width and h < self.height: + top = random.randint(0, self.height - h) + left = random.randint(0, self.width - w) + + num_masked = mask[top : top + h, left : left + w].sum() + # Overlap + if 0 < h * w - num_masked <= max_mask_patches: + for i in range(top, top + h): + for j in range(left, left + w): + if mask[i, j] == 0: + mask[i, j] = 1 + delta += 1 + + if delta > 0: + break + return delta + + def __call__(self): + mask = np.zeros(shape=self.get_shape(), dtype=int) + mask_count = 0 + while mask_count < self.num_masking_patches: + max_mask_patches = self.num_masking_patches - mask_count + max_mask_patches = min(max_mask_patches, self.max_num_patches) + + delta = self._mask(mask, max_mask_patches) + if delta == 0: + break + else: + mask_count += delta + + return mask + + +""" +This implementation is based on +https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py +Licensed under The MIT License +""" + + +class MaskingGenerator3D: + def __init__( + self, + mask_window_size, + num_masking_patches, + min_num_patches=16, + max_num_patches=None, + min_aspect=0.3, + max_aspect=None, + ): + self.temporal, self.height, self.width = mask_window_size + self.num_masking_patches = num_masking_patches + self.min_num_patches = min_num_patches + self.max_num_patches = ( + num_masking_patches if max_num_patches is None else max_num_patches + ) + max_aspect = max_aspect or 1 / min_aspect + self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) + + def __repr__(self): + repr_str = "Generator(%d, %d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % ( + self.temporal, + self.height, + self.width, + self.min_num_patches, + self.max_num_patches, + self.num_masking_patches, + self.log_aspect_ratio[0], + self.log_aspect_ratio[1], + ) + return repr_str + + def get_shape(self): + return self.temporal, self.height, self.width + + def _mask(self, mask, max_mask_patches): + delta = 0 + for _ in range(100): + target_area = random.uniform(self.min_num_patches, self.max_num_patches) + aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) + h = int(round(math.sqrt(target_area * aspect_ratio))) + w = int(round(math.sqrt(target_area / aspect_ratio))) + t = random.randint(1, self.temporal) # ! + if w < self.width and h < self.height: + top = random.randint(0, self.height - h) + left = random.randint(0, self.width - w) + front = random.randint(0, self.temporal - t) + + num_masked = mask[ + front : front + t, top : top + h, left : left + w + ].sum() + # Overlap + if 0 < h * w * t - num_masked <= max_mask_patches: + for i in range(front, front + t): + for j in range(top, top + h): + for k in range(left, left + w): + if mask[i, j, k] == 0: + mask[i, j, k] = 1 + delta += 1 + + if delta > 0: + break + return delta + + def __call__(self): + mask = np.zeros(shape=self.get_shape(), dtype=int) + mask_count = 0 + while mask_count < self.num_masking_patches: + max_mask_patches = self.num_masking_patches - mask_count + + delta = self._mask(mask, max_mask_patches) + if delta == 0: + break + else: + mask_count += delta + + return mask + + def transforms_imagenet_train( img_size=224, scale=None, @@ -830,9 +976,7 @@ def transforms_imagenet_train( img_size = img_size scale = tuple(scale or (0.08, 1.0)) # default imagenet scale range - ratio = tuple( - ratio or (3.0 / 4.0, 4.0 / 3.0) - ) # default imagenet ratio range + ratio = tuple(ratio or (3.0 / 4.0, 4.0 / 3.0)) # default imagenet ratio range primary_tfl = [ RandomResizedCropAndInterpolation( img_size, scale=scale, ratio=ratio, interpolation=interpolation @@ -980,9 +1124,7 @@ def color_jitter_video_ssl( def augment_raw_frames(frames, time_diff_prob=0.0, gaussian_prob=0.0): frames = frames.float() if gaussian_prob > 0.0: - blur_trans = tv.transforms.RandomApply( - [GaussianBlurVideo()], p=gaussian_prob - ) + blur_trans = tv.transforms.RandomApply([GaussianBlurVideo()], p=gaussian_prob) frames = blur_trans(frames) time_diff_out = False @@ -1000,7 +1142,7 @@ def augment_raw_frames(frames, time_diff_prob=0.0, gaussian_prob=0.0): return frames, time_diff_out -class GaussianBlur(object): +class GaussianBlur: """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709""" def __init__(self, sigma=[0.1, 2.0]): @@ -1015,10 +1157,8 @@ def __call__(self, x): return x -class GaussianBlurVideo(object): - def __init__( - self, sigma_min=[0.0, 0.1], sigma_max=[0.0, 2.0], use_PIL=False - ): +class GaussianBlurVideo: + def __init__(self, sigma_min=[0.0, 0.1], sigma_max=[0.0, 2.0], use_PIL=False): self.sigma_min = sigma_min self.sigma_max = sigma_max diff --git a/slowfast/datasets/utils.py b/slowfast/datasets/utils.py index df8222975..05257f025 100644 --- a/slowfast/datasets/utils.py +++ b/slowfast/datasets/utils.py @@ -1,18 +1,22 @@ #!/usr/bin/env python3 import logging -import numpy as np import os import random import time from collections import defaultdict + import cv2 +import numpy as np import torch -from torch.utils.data.distributed import DistributedSampler - from slowfast.utils.env import pathmgr +from torch.utils.data.distributed import DistributedSampler +from torchvision import transforms from . import transform as transform +from .random_erasing import RandomErasing +from .transform import create_random_augment + logger = logging.getLogger(__name__) @@ -259,9 +263,7 @@ def load_image_lists(frame_list_file, prefix="", return_list=False): image_paths[video_name].append(path) frame_labels = row[-1].replace('"', "") if frame_labels != "": - labels[video_name].append( - [int(x) for x in frame_labels.split(",")] - ) + labels[video_name].append([int(x) for x in frame_labels.split(",")]) else: labels[video_name].append([]) @@ -348,3 +350,95 @@ def loader_worker_init_fn(dataset): dataset (torch.utils.data.Dataset): the given dataset. """ return None + + +def aug_frame( + cfg, + mode, + rand_erase, + frames, + spatial_sample_index, + min_scale, + max_scale, + crop_size, +): + """ + Perform augmentations on the given video frames, including + random augmentation, normalization, spatial sampling and optional random + erasing. + Args: + cfg (CfgNode): configs. + mode (string): Options includes `train`, `val`, or `test` mode. + rand_erase (bool): if performing random erasing. + frames (tensor): frames of images sampled from the video. The + dimension is `num frames` x `height` x `width` x `channel`. + spatial_sample_index (int): if -1, perform random spatial sampling. + If 0, 1, or 2, perform left, center, right crop if width is larger + thanheight, and perform top, center, buttom crop if height is larger + than width. + min_scale (int): the minimal size of scaling. + max_scale (int): the maximal size of scaling. + crop_size (int): the size of height and width used to crop the + frames. + Returns: + frames (tensor): spatially sampled frames. + """ + if cfg.AUG.AA_TYPE: + aug_transform = create_random_augment( + input_size=(frames.size(1), frames.size(2)), + auto_augment=cfg.AUG.AA_TYPE, + interpolation=cfg.AUG.INTERPOLATION, + ) + # T H W C -> T C H W. + frames = frames.permute(0, 3, 1, 2) + list_img = _frame_to_list_img(frames) + list_img = aug_transform(list_img) + frames = _list_img_to_frames(list_img) + frames = frames.permute(0, 2, 3, 1) + + frames = tensor_normalize(frames, cfg.DATA.MEAN, cfg.DATA.STD) + # T H W C -> C T H W. + frames = frames.permute(3, 0, 1, 2) + # Perform data augmentation. + scl, asp = ( + cfg.DATA.TRAIN_JITTER_SCALES_RELATIVE, + cfg.DATA.TRAIN_JITTER_ASPECT_RELATIVE, + ) + relative_scales = None if (mode not in ["train"] or len(scl) == 0) else scl + relative_aspect = None if (mode not in ["train"] or len(asp) == 0) else asp + frames = spatial_sampling( + frames, + spatial_idx=spatial_sample_index, + min_scale=min_scale, + max_scale=max_scale, + crop_size=crop_size, + random_horizontal_flip=cfg.DATA.RANDOM_FLIP, + inverse_uniform_sampling=cfg.DATA.INV_UNIFORM_SAMPLE, + aspect_ratio=relative_aspect, + scale=relative_scales, + motion_shift=cfg.DATA.TRAIN_JITTER_MOTION_SHIFT if mode in ["train"] else False, + ) + + if rand_erase: + erase_transform = RandomErasing( + cfg.AUG.RE_PROB, + mode=cfg.AUG.RE_MODE, + max_count=cfg.AUG.RE_COUNT, + num_splits=cfg.AUG.RE_COUNT, + device="cpu", + ) + frames = frames.permute(1, 0, 2, 3) + frames = erase_transform(frames) + frames = frames.permute(1, 0, 2, 3) + + return frames + + +def _frame_to_list_img(frames): + img_list = [transforms.ToPILImage()(frames[i]) for i in range(frames.size(0))] + return img_list + + +def _list_img_to_frames(img_list): + img_list = [transforms.ToTensor()(img) for img in img_list] + return torch.stack(img_list) diff --git a/slowfast/models/__init__.py b/slowfast/models/__init__.py index 88e806198..88de25564 100644 --- a/slowfast/models/__init__.py +++ b/slowfast/models/__init__.py @@ -1,18 +1,18 @@ #!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -from .build import MODEL_REGISTRY, build_model # noqa +from .build import build_model, MODEL_REGISTRY # noqa from .contrastive import ContrastiveModel # noqa -from .custom_video_model_builder import * # noqa -from .video_model_builder import ResNet, SlowFast # noqa +from .masked import MaskMViT # noqa +from .video_model_builder import MViT, ResNet, SlowFast # noqa try: from .ptv_model_builder import ( PTVCSN, - PTVX3D, PTVR2plus1D, PTVResNet, PTVSlowFast, + PTVX3D, ) # noqa except Exception: print("Please update your PyTorchVideo to latest master") diff --git a/slowfast/models/attention.py b/slowfast/models/attention.py index 22abc77a9..29c80f046 100644 --- a/slowfast/models/attention.py +++ b/slowfast/models/attention.py @@ -6,9 +6,8 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.nn.init import trunc_normal_ - from slowfast.models.common import DropPath, Mlp +from torch.nn.init import trunc_normal_ def attention_pool(tensor, pool, thw_shape, has_cls_embed=True, norm=None): @@ -27,9 +26,7 @@ def attention_pool(tensor, pool, thw_shape, has_cls_embed=True, norm=None): B, N, L, C = tensor.shape T, H, W = thw_shape - tensor = ( - tensor.reshape(B * N, T, H, W, C).permute(0, 4, 1, 2, 3).contiguous() - ) + tensor = tensor.reshape(B * N, T, H, W, C).permute(0, 4, 1, 2, 3).contiguous() tensor = pool(tensor) @@ -80,15 +77,13 @@ def cal_rel_pos_spatial( q_h_ratio = max(k_h / q_h, 1.0) k_h_ratio = max(q_h / k_h, 1.0) dist_h = ( - torch.arange(q_h)[:, None] * q_h_ratio - - torch.arange(k_h)[None, :] * k_h_ratio + torch.arange(q_h)[:, None] * q_h_ratio - torch.arange(k_h)[None, :] * k_h_ratio ) dist_h += (k_h - 1) * k_h_ratio q_w_ratio = max(k_w / q_w, 1.0) k_w_ratio = max(q_w / k_w, 1.0) dist_w = ( - torch.arange(q_w)[:, None] * q_w_ratio - - torch.arange(k_w)[None, :] * k_w_ratio + torch.arange(q_w)[:, None] * q_w_ratio - torch.arange(k_w)[None, :] * k_w_ratio ) dist_w += (k_w - 1) * k_w_ratio @@ -101,12 +96,8 @@ def cal_rel_pos_spatial( B, n_head, q_N, dim = q.shape r_q = q[:, :, sp_idx:].reshape(B, n_head, q_t, q_h, q_w, dim) - rel_h_q = torch.einsum( - "bythwc,hkc->bythwk", r_q, Rh - ) # [B, H, q_t, qh, qw, k_h] - rel_w_q = torch.einsum( - "bythwc,wkc->bythwk", r_q, Rw - ) # [B, H, q_t, qh, qw, k_w] + rel_h_q = torch.einsum("bythwc,hkc->bythwk", r_q, Rh) # [B, H, q_t, qh, qw, k_h] + rel_w_q = torch.einsum("bythwc,wkc->bythwk", r_q, Rw) # [B, H, q_t, qh, qw, k_w] attn[:, :, sp_idx:, sp_idx:] = ( attn[:, :, sp_idx:, sp_idx:].view(B, -1, q_t, q_h, q_w, k_t, k_h, k_w) @@ -132,8 +123,7 @@ def cal_rel_pos_temporal(attn, q, has_cls_embed, q_shape, k_shape, rel_pos_t): q_t_ratio = max(k_t / q_t, 1.0) k_t_ratio = max(q_t / k_t, 1.0) dist_t = ( - torch.arange(q_t)[:, None] * q_t_ratio - - torch.arange(k_t)[None, :] * k_t_ratio + torch.arange(q_t)[:, None] * q_t_ratio - torch.arange(k_t)[None, :] * k_t_ratio ) dist_t += (k_t - 1) * k_t_ratio Rt = rel_pos_t[dist_t.long()] @@ -142,9 +132,7 @@ def cal_rel_pos_temporal(attn, q, has_cls_embed, q_shape, k_shape, rel_pos_t): r_q = q[:, :, sp_idx:].reshape(B, n_head, q_t, q_h, q_w, dim) # [B, H, q_t, q_h, q_w, dim] -> [q_t, B, H, q_h, q_w, dim] -> [q_t, B*H*q_h*q_w, dim] - r_q = r_q.permute(2, 0, 1, 3, 4, 5).reshape( - q_t, B * n_head * q_h * q_w, dim - ) + r_q = r_q.permute(2, 0, 1, 3, 4, 5).reshape(q_t, B * n_head * q_h * q_w, dim) # [q_t, B*H*q_h*q_w, dim] * [q_t, dim, k_t] = [q_t, B*H*q_h*q_w, k_t] -> [B*H*q_h*q_w, q_t, k_t] rel = torch.matmul(r_q, Rt.transpose(1, 2)).transpose(0, 1) @@ -193,6 +181,7 @@ def __init__( head_dim = dim_out // num_heads self.scale = head_dim**-0.5 self.has_cls_embed = has_cls_embed + self.mode = mode padding_q = [int(q // 2) for q in kernel_q] padding_kv = [int(kv // 2) for kv in kernel_kv] @@ -212,7 +201,6 @@ def __init__( kernel_q = () if numpy.prod(kernel_kv) == 1 and numpy.prod(stride_kv) == 1: kernel_kv = () - self.mode = mode if mode in ("avg", "max"): pool_op = nn.MaxPool3d if mode == "max" else nn.AvgPool3d @@ -296,9 +284,7 @@ def __init__( trunc_normal_(self.rel_pos_h, std=0.02) trunc_normal_(self.rel_pos_w, std=0.02) if self.rel_pos_temporal: - self.rel_pos_t = nn.Parameter( - torch.zeros(2 * input_size[0] - 1, head_dim) - ) + self.rel_pos_t = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) if not rel_pos_zero_init: trunc_normal_(self.rel_pos_t, std=0.02) @@ -325,81 +311,45 @@ def forward(self, x, thw_shape): q, k, v = qkv[0], qkv[1], qkv[2] else: q = k = v = x - q = ( - self.q(q) - .reshape(B, N, self.num_heads, -1) - .permute(0, 2, 1, 3) - ) - k = ( - self.k(k) - .reshape(B, N, self.num_heads, -1) - .permute(0, 2, 1, 3) - ) - v = ( - self.v(v) - .reshape(B, N, self.num_heads, -1) - .permute(0, 2, 1, 3) - ) + q = self.q(q).reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) + k = self.k(k).reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) + v = self.v(v).reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) q, q_shape = attention_pool( q, self.pool_q, thw_shape, has_cls_embed=self.has_cls_embed, - norm=self.norm_q if hasattr(self, "norm_q") else None, + norm=getattr(self, "norm_q", None), ) k, k_shape = attention_pool( k, self.pool_k, thw_shape, has_cls_embed=self.has_cls_embed, - norm=self.norm_k if hasattr(self, "norm_k") else None, + norm=getattr(self, "norm_k", None), ) v, v_shape = attention_pool( v, self.pool_v, thw_shape, has_cls_embed=self.has_cls_embed, - norm=self.norm_v if hasattr(self, "norm_v") else None, + norm=getattr(self, "norm_v", None), ) if self.pool_first: - q_N = ( - numpy.prod(q_shape) + 1 - if self.has_cls_embed - else numpy.prod(q_shape) - ) - k_N = ( - numpy.prod(k_shape) + 1 - if self.has_cls_embed - else numpy.prod(k_shape) - ) - v_N = ( - numpy.prod(v_shape) + 1 - if self.has_cls_embed - else numpy.prod(v_shape) - ) + q_N = numpy.prod(q_shape) + 1 if self.has_cls_embed else numpy.prod(q_shape) + k_N = numpy.prod(k_shape) + 1 if self.has_cls_embed else numpy.prod(k_shape) + v_N = numpy.prod(v_shape) + 1 if self.has_cls_embed else numpy.prod(v_shape) q = q.permute(0, 2, 1, 3).reshape(B, q_N, -1) - q = ( - self.q(q) - .reshape(B, q_N, self.num_heads, -1) - .permute(0, 2, 1, 3) - ) + q = self.q(q).reshape(B, q_N, self.num_heads, -1).permute(0, 2, 1, 3) v = v.permute(0, 2, 1, 3).reshape(B, v_N, -1) - v = ( - self.v(v) - .reshape(B, v_N, self.num_heads, -1) - .permute(0, 2, 1, 3) - ) + v = self.v(v).reshape(B, v_N, self.num_heads, -1).permute(0, 2, 1, 3) k = k.permute(0, 2, 1, 3).reshape(B, k_N, -1) - k = ( - self.k(k) - .reshape(B, k_N, self.num_heads, -1) - .permute(0, 2, 1, 3) - ) + k = self.k(k).reshape(B, k_N, self.num_heads, -1).permute(0, 2, 1, 3) N = q.shape[2] attn = (q * self.scale) @ k.transpose(-2, -1) @@ -454,6 +404,7 @@ def __init__( qk_scale=None, drop_rate=0.0, drop_path=0.0, + layer_scale_init_value=0.0, act_layer=nn.GELU, norm_layer=nn.LayerNorm, up_rate=None, @@ -501,9 +452,7 @@ def __init__( residual_pooling=residual_pooling, separate_qkv=separate_qkv, ) - self.drop_path = ( - DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(att_dim) mlp_hidden_dim = int(att_dim * mlp_ratio) self.has_cls_embed = has_cls_embed @@ -519,18 +468,27 @@ def __init__( act_layer=act_layer, drop_rate=drop_rate, ) + if layer_scale_init_value > 0: + self.gamma_1 = nn.Parameter( + layer_scale_init_value * torch.ones((dim)), requires_grad=True + ) + self.gamma_2 = nn.Parameter( + layer_scale_init_value * torch.ones((dim_out)), + requires_grad=True, + ) + else: + self.gamma_1, self.gamma_2 = None, None + if dim != dim_out: self.proj = nn.Linear(dim, dim_out) self.pool_skip = ( - nn.MaxPool3d( - kernel_skip, stride_skip, padding_skip, ceil_mode=False - ) - if len(kernel_skip) > 0 + nn.MaxPool3d(kernel_skip, stride_skip, padding_skip, ceil_mode=False) + if len(stride_skip) > 0 and numpy.prod(stride_skip) > 1 else None ) - def forward(self, x, thw_shape): + def forward(self, x, thw_shape=None): x_norm = self.norm1(x) x_block, thw_shape_new = self.attn(x_norm, thw_shape) if self.dim_mul_in_att and self.dim != self.dim_out: @@ -538,10 +496,19 @@ def forward(self, x, thw_shape): x_res, _ = attention_pool( x, self.pool_skip, thw_shape, has_cls_embed=self.has_cls_embed ) - x = x_res + self.drop_path(x_block) + if self.gamma_1 is not None: + x = x_res + self.drop_path(self.gamma_1 * x_block) + else: + x = x_res + self.drop_path(x_block) x_norm = self.norm2(x) x_mlp = self.mlp(x_norm) if not self.dim_mul_in_att and self.dim != self.dim_out: x = self.proj(x_norm) - x = x + self.drop_path(x_mlp) - return x, thw_shape_new + if self.gamma_2 is not None: + x = x + self.drop_path(self.gamma_2 * x_mlp) + else: + x = x + self.drop_path(x_mlp) + if thw_shape: + return x, thw_shape_new + else: + return x diff --git a/slowfast/models/batchnorm_helper.py b/slowfast/models/batchnorm_helper.py index 1351d0a65..b020eccd8 100644 --- a/slowfast/models/batchnorm_helper.py +++ b/slowfast/models/batchnorm_helper.py @@ -4,13 +4,13 @@ """BatchNorm (BN) utility functions and custom batch-size BN implementations""" from functools import partial + import torch import torch.nn as nn - -from pytorchvideo.layers.batch_norm import ( +from pytorchvideo.layers.batch_norm import ( # noqa: F401 NaiveSyncBatchNorm1d, NaiveSyncBatchNorm3d, -) # noqa +) def get_norm(cfg): diff --git a/slowfast/models/build.py b/slowfast/models/build.py index e68feae69..6d35ba29f 100644 --- a/slowfast/models/build.py +++ b/slowfast/models/build.py @@ -3,13 +3,10 @@ """Model construction functions.""" +import slowfast.utils.logging as logging import torch from fvcore.common.registry import Registry -from torch.distributed.algorithms.ddp_comm_hooks import ( - default as comm_hooks_default, -) - -import slowfast.utils.logging as logging +from torch.distributed.algorithms.ddp_comm_hooks import default as comm_hooks_default logger = logging.get_logger(__name__) @@ -31,13 +28,13 @@ def build_model(cfg, gpu_id=None): gpu_id (Optional[int]): specify the gpu index to build model. """ if torch.cuda.is_available(): - assert ( - cfg.NUM_GPUS <= torch.cuda.device_count() - ), "Cannot use more GPU devices than available" + assert cfg.NUM_GPUS <= torch.cuda.device_count(), ( + "Cannot use more GPU devices than available" + ) else: - assert ( - cfg.NUM_GPUS == 0 - ), "Cuda is not available. Please set `NUM_GPUS: 0 for running on CPUs." + assert cfg.NUM_GPUS == 0, ( + "Cuda is not available. Please set `NUM_GPUS: 0 for running on CPUs." + ) # Construct the model name = cfg.MODEL.MODEL_NAME @@ -53,9 +50,7 @@ def build_model(cfg, gpu_id=None): process_group = apex.parallel.create_syncbn_process_group( group_size=cfg.BN.NUM_SYNC_DEVICES ) - model = apex.parallel.convert_syncbn_model( - model, process_group=process_group - ) + model = apex.parallel.convert_syncbn_model(model, process_group=process_group) if cfg.NUM_GPUS: if gpu_id is None: @@ -72,10 +67,12 @@ def build_model(cfg, gpu_id=None): module=model, device_ids=[cur_device], output_device=cur_device, - find_unused_parameters=True - if cfg.MODEL.DETACH_FINAL_FC - or cfg.MODEL.MODEL_NAME == "ContrastiveModel" - else False, + find_unused_parameters=( + True + if cfg.MODEL.DETACH_FINAL_FC + or cfg.MODEL.MODEL_NAME == "ContrastiveModel" + else False + ), ) if cfg.MODEL.FP16_ALLREDUCE: model.register_comm_hook( diff --git a/slowfast/models/common.py b/slowfast/models/common.py index 79fd494f9..2abaf256f 100644 --- a/slowfast/models/common.py +++ b/slowfast/models/common.py @@ -68,3 +68,79 @@ def __init__(self, drop_prob=None): def forward(self, x): return drop_path(x, self.drop_prob, self.training) + + +class TwoStreamFusion(nn.Module): + def __init__(self, mode, dim=None, kernel=3, padding=1): + """ + A general constructor for neural modules fusing two equal sized tensors + in forward. Following options are supported: + + "add" / "max" / "min" / "avg" : respective operations on the two halves. + "concat" : NOOP. + "concat_linear_{dim_mult}_{drop_rate}" : MLP to fuse with hidden dim "dim_mult" + (optional, def 1.) higher than input dim + with optional dropout "drop_rate" (def: 0.) + "ln+concat_linear_{dim_mult}_{drop_rate}" : perform MLP after layernorm on the input. + + """ + super().__init__() + self.mode = mode + if mode == "add": + self.fuse_fn = lambda x: torch.stack(torch.chunk(x, 2, dim=2)).sum(dim=0) + elif mode == "max": + self.fuse_fn = ( + lambda x: torch.stack(torch.chunk(x, 2, dim=2)).max(dim=0).values + ) + elif mode == "min": + self.fuse_fn = ( + lambda x: torch.stack(torch.chunk(x, 2, dim=2)).min(dim=0).values + ) + elif mode == "avg": + self.fuse_fn = lambda x: torch.stack(torch.chunk(x, 2, dim=2)).mean(dim=0) + elif mode == "concat": + # x itself is the channel concat version + self.fuse_fn = lambda x: x + elif "concat_linear" in mode: + if len(mode.split("_")) == 2: + dim_mult = 1.0 + drop_rate = 0.0 + elif len(mode.split("_")) == 3: + dim_mult = float(mode.split("_")[-1]) + drop_rate = 0.0 + + elif len(mode.split("_")) == 4: + dim_mult = float(mode.split("_")[-2]) + drop_rate = float(mode.split("_")[-1]) + else: + raise NotImplementedError + + if mode.split("+")[0] == "ln": + self.fuse_fn = nn.Sequential( + nn.LayerNorm(dim), + Mlp( + in_features=dim, + hidden_features=int(dim * dim_mult), + act_layer=nn.GELU, + out_features=dim, + drop_rate=drop_rate, + ), + ) + else: + self.fuse_fn = Mlp( + in_features=dim, + hidden_features=int(dim * dim_mult), + act_layer=nn.GELU, + out_features=dim, + drop_rate=drop_rate, + ) + + else: + raise NotImplementedError + + def forward(self, x): + if "concat_linear" in self.mode: + return self.fuse_fn(x) + x + + else: + return self.fuse_fn(x) diff --git a/slowfast/models/contrastive.py b/slowfast/models/contrastive.py index 516e783b6..08ac70c0d 100644 --- a/slowfast/models/contrastive.py +++ b/slowfast/models/contrastive.py @@ -1,18 +1,16 @@ #!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -"""CSC model.""" - import math -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F +import numpy as np import slowfast.models.losses as losses import slowfast.utils.distributed as du import slowfast.utils.logging as logging -from slowfast.models.video_model_builder import X3D, MViT, ResNet, SlowFast +import torch +import torch.nn as nn +import torch.nn.functional as F +from slowfast.models.video_model_builder import MViT, ResNet, SlowFast, X3D from .build import MODEL_REGISTRY @@ -56,21 +54,16 @@ def __init__(self, cfg): self.train_labels = np.zeros((0,), dtype=np.int32) self.num_pos = 2 self.num_crops = ( - self.cfg.DATA.TRAIN_CROP_NUM_TEMPORAL - * self.cfg.DATA.TRAIN_CROP_NUM_SPATIAL - ) - self.nce_loss_fun = losses.get_loss_func("contrastive_loss")( - reduction="mean" + self.cfg.DATA.TRAIN_CROP_NUM_TEMPORAL * self.cfg.DATA.TRAIN_CROP_NUM_SPATIAL ) + self.nce_loss_fun = losses.get_loss_func("contrastive_loss")(reduction="mean") assert self.cfg.MODEL.LOSS_FUNC == "contrastive_loss" self.softmax = nn.Softmax(dim=1).cuda() if self.type == "mem": self.mem_type = cfg.CONTRASTIVE.MEM_TYPE if self.mem_type == "1d": - self.memory = Memory1D( - self.length, self.duration, self.dim, cfg - ) + self.memory = Memory1D(self.length, self.duration, self.dim, cfg) else: self.memory = Memory(self.length, self.duration, self.dim, cfg) self.examplar_type = "video" @@ -117,8 +110,7 @@ def __init__(self, cfg): "queue_swav", torch.zeros( 2, # = args.crops_for_assign - self.cfg.CONTRASTIVE.SWAV_QEUE_LEN - // du.get_world_size(), + self.cfg.CONTRASTIVE.SWAV_QEUE_LEN // du.get_world_size(), self.dim, ), ) @@ -255,10 +247,7 @@ def momentum_anneal_cosine(self, epoch_exact): self.mmt = ( 1 - (1 - self.cfg.CONTRASTIVE.MOMENTUM) - * ( - math.cos(math.pi * epoch_exact / self.cfg.SOLVER.MAX_EPOCH) - + 1.0 - ) + * (math.cos(math.pi * epoch_exact / self.cfg.SOLVER.MAX_EPOCH) + 1.0) * 0.5 ) @@ -270,9 +259,7 @@ def _dequeue_and_enqueue(self, keys, extra_keys=None): ): # TODO: add multiview negatives keys_queue_update = [keys[0]] else: - assert ( - len(keys) > 0 - ), "need to have multiple views for adding them to queue" + assert len(keys) > 0, "need to have multiple views for adding them to queue" keys_queue_update = [] keys_queue_update += keys if extra_keys: @@ -301,9 +288,7 @@ def batch_clips(self, clips): if i == 0: clips_batched[j] = view else: - clips_batched[j] = torch.cat( - [clips_batched[j], view], dim=0 - ) + clips_batched[j] = torch.cat([clips_batched[j], view], dim=0) del view return clips_batched @@ -348,9 +333,7 @@ def compute_key_feat( for tk in hist_time: tk = self.l2_norm(tk) if self._batch_shuffle_on: - tk = self._batch_unshuffle( - tk, idx_restore - ).detach() + tk = self._batch_unshuffle(tk, idx_restore).detach() tks.append(tk) pred_keys.append(tks) x_hist = self.l2_norm(hist_feat) @@ -372,9 +355,7 @@ def compute_key_feat( else: return keys - def forward( - self, clips, index=None, time=None, epoch_exact=None, keys=None - ): + def forward(self, clips, index=None, time=None, epoch_exact=None, keys=None): if epoch_exact is not None and self.momentum_annealing: self.momentum_anneal_cosine(epoch_exact) @@ -416,9 +397,7 @@ def forward( ), ).cuda() else: - time_ind = torch.zeros( - size=(batch_size, self.k + 1), dtype=int - ).cuda() + time_ind = torch.zeros(size=(batch_size, self.k + 1), dtype=int).cuda() if self.examplar_type == "clip": # Diff clip from same video are negative. @@ -477,17 +456,13 @@ def forward( return self.eval_knn(q_knn) if keys is None: - keys = self.compute_key_feat( - clips_k, compute_predictor_keys=False - ) + keys = self.compute_key_feat(clips_k, compute_predictor_keys=False) auto_enqueue_keys = True else: auto_enqueue_keys = False # score computation - queue_neg = torch.einsum( - "nc,kc->nk", [q, self.queue_x.clone().detach()] - ) + queue_neg = torch.einsum("nc,kc->nk", [q, self.queue_x.clone().detach()]) for k, key in enumerate(keys): out_pos = torch.einsum("nc,nc->n", [q, key]).unsqueeze(-1) @@ -547,15 +522,11 @@ def forward( if not self.training: return self.eval_knn(q_knn) - ind_clips = np.arange( - n_clips - ) # clips come ordered temporally from decoder + ind_clips = np.arange(n_clips) # clips come ordered temporally from decoder # rest down is for training if keys is None: - keys = self.compute_key_feat( - clips_key, compute_predictor_keys=False - ) + keys = self.compute_key_feat(clips_key, compute_predictor_keys=False) if self.cfg.CONTRASTIVE.SEQUENTIAL: loss_reg = self.sim_loss(predictors[0], keys[0]) @@ -586,8 +557,7 @@ def forward( # loss_pos = self.sim_loss(q1_proj, q2_proj) dummy_logits = torch.cat( ( - 9999.0 - * torch.ones((len(index), 1), dtype=torch.float).cuda(), + 9999.0 * torch.ones((len(index), 1), dtype=torch.float).cuda(), torch.zeros((len(index), self.k), dtype=torch.float).cuda(), ), dim=1, @@ -608,16 +578,14 @@ def forward( if not self.training: return self.eval_knn(proj_1) n_clips = len(clips) - ind_clips = np.arange( - n_clips - ) # clips come ordered temporally from decoder + ind_clips = np.arange(n_clips) # clips come ordered temporally from decoder clip_q = clips[0] if self.swav_use_public_code: # uses official code of SwAV from # https://github.com/facebookresearch/swav/blob/master/main_swav.py with torch.no_grad(): - m = self.module if hasattr(self, "module") else self + m = getattr(self, "module", self) w = m.swav_prototypes.weight.data.clone() w = nn.functional.normalize(w, dim=1, p=2) m.swav_prototypes.weight.copy_(w) @@ -625,7 +593,7 @@ def forward( bs = clips[0][0].size(0) output, embedding = [], [] - for i, clip_q in enumerate(clips): + for clip_q in clips: x = self.run_swav_orig_encoder_q(clip_q) embedding.append(x[0]) output.append(x[1]) @@ -635,9 +603,7 @@ def forward( loss_swav = 0 swav_extra_crops = n_clips - 2 - self.swav_crops_for_assign = np.arange( - n_clips - swav_extra_crops - ) + self.swav_crops_for_assign = np.arange(n_clips - swav_extra_crops) for i, crop_id in enumerate(self.swav_crops_for_assign): with torch.no_grad(): out = output[bs * crop_id : bs * (crop_id + 1)] @@ -658,9 +624,7 @@ def forward( out, ) ) - self.queue_swav[i, bs:] = self.queue_swav[ - i, :-bs - ].clone() + self.queue_swav[i, bs:] = self.queue_swav[i, :-bs].clone() self.queue_swav[i, :bs] = embedding[ crop_id * bs : (crop_id + 1) * bs ] @@ -674,9 +638,7 @@ def forward( subloss = 0 for v in np.delete(np.arange(n_clips), crop_id): p = self.softmax(output[bs * v : bs * (v + 1)] / self.T) - subloss -= torch.mean( - torch.sum(q * torch.log(p), dim=1) - ) + subloss -= torch.mean(torch.sum(q * torch.log(p), dim=1)) loss_swav += subloss / (n_clips - 1) loss_swav /= len(self.swav_crops_for_assign) else: @@ -695,17 +657,13 @@ def forward( ).detach() out_1 = torch.cat( ( - torch.mm( - self.queue_swav[0].detach(), swav_prototypes - ), + torch.mm(self.queue_swav[0].detach(), swav_prototypes), out_1, ) ) out_2 = torch.cat( ( - torch.mm( - self.queue_swav[1].detach(), swav_prototypes - ), + torch.mm(self.queue_swav[1].detach(), swav_prototypes), out_2, ) ) @@ -724,8 +682,7 @@ def forward( self.knn_mem_update(q_knn, index) dummy_logits = torch.cat( ( - 9999.0 - * torch.ones((len(index), 1), dtype=torch.float).cuda(), + 9999.0 * torch.ones((len(index), 1), dtype=torch.float).cuda(), torch.zeros((len(index), self.k), dtype=torch.float).cuda(), ), dim=1, @@ -759,9 +716,7 @@ def forward( for loss_id in range(len(self.pos_mask)): pos = torch.sum(similarity * self.pos_mask[loss_id], 1) neg = torch.sum(similarity * self.neg_mask, 1) - idx = ( - 1 - torch.sum(self.pos_mask[loss_id], 1) > 0 - ).detach() + idx = (1 - torch.sum(self.pos_mask[loss_id], 1) > 0).detach() term_prob = pos / (pos + neg) term_prob[idx] = 1.0 term_loss = torch.log(term_prob) @@ -775,18 +730,14 @@ def forward( q2 = du.AllGatherWithGradient.apply(q2) out = torch.cat([q, q2], dim=0) # [2*B, 2*B] - sim_matrix = torch.exp( - torch.mm(out, out.t().contiguous()) / self.T - ) + sim_matrix = torch.exp(torch.mm(out, out.t().contiguous()) / self.T) # SANITY: mask = ( torch.ones_like(sim_matrix) - torch.eye(out.shape[0], device=sim_matrix.device) ).bool() # [2*B, 2*B-1] - sim_matrix = sim_matrix.masked_select(mask).view( - out.shape[0], -1 - ) + sim_matrix = sim_matrix.masked_select(mask).view(out.shape[0], -1) # compute loss pos_sim = torch.exp(torch.sum(q * q2, dim=-1) / self.T) # [2*B] @@ -795,8 +746,7 @@ def forward( self.knn_mem_update(q_knn, index) dummy_logits = torch.cat( ( - 9999.0 - * torch.ones((len(index), 1), dtype=torch.float).cuda(), + 9999.0 * torch.ones((len(index), 1), dtype=torch.float).cuda(), torch.zeros((len(index), self.k), dtype=torch.float).cuda(), ), dim=1, @@ -929,9 +879,7 @@ def __init__(self, power=2, dim=1): self.power = power def forward(self, x): - norm = ( - x.pow(self.power).sum(self.dim, keepdim=True).pow(1.0 / self.power) - ) + norm = x.pow(self.power).sum(self.dim, keepdim=True).pow(1.0 / self.power) out = x.div(norm) return out @@ -983,9 +931,7 @@ def get(self, ind, time, interp=False): selected_mem = mem_t0 * (1 - w_t1) + mem_t1 * w_t1 else: # logger.info("1dmem get ind shape {} time shape {}".format(ind.shape, time.shape)) - selected_mem = self.memory[ - ind.view(-1), time.long().view(-1), : - ] + selected_mem = self.memory[ind.view(-1), time.long().view(-1), :] out = selected_mem.view(batch_size, -1, self.dim) return out @@ -1009,12 +955,12 @@ def update(self, mem, momentum, ind, time, interp=False): # mem = mem.squeeze() duo_update = False if duo_update: - update_t0 = ( - mem * w_t0 + mem_t0 * w_t1 - ) * momentum + mem_t0 * (1 - momentum) - update_t1 = ( - mem * w_t1 + mem_t1 * w_t0 - ) * momentum + mem_t1 * (1 - momentum) + update_t0 = (mem * w_t0 + mem_t0 * w_t1) * momentum + mem_t0 * ( + 1 - momentum + ) + update_t1 = (mem * w_t1 + mem_t1 * w_t0) * momentum + mem_t1 * ( + 1 - momentum + ) else: update_t0 = mem * w_t0 * momentum + mem_t0 * (1 - momentum) update_t1 = mem * w_t1 * momentum + mem_t1 * (1 - momentum) @@ -1032,9 +978,9 @@ def update(self, mem, momentum, ind, time, interp=False): # logger.info("1dmem set ind shape {} time shape {}".format(ind.shape, time.shape)) # my version - self.memory[ - ind.view(-1), time.long().view(-1), : - ] = mem_update.squeeze() + self.memory[ind.view(-1), time.long().view(-1), :] = ( + mem_update.squeeze() + ) return def forward(self, inputs): @@ -1082,7 +1028,7 @@ def update(self, mem, momentum, ind, time, interp=False): return -def cancel_swav_gradients(model, cfg, epoch_exact): +def contrastive_parameter_surgery(model, cfg, epoch_exact, cur_iter): # cancel some gradients in first epoch of SwAV if ( cfg.MODEL.MODEL_NAME == "ContrastiveModel" @@ -1092,4 +1038,63 @@ def cancel_swav_gradients(model, cfg, epoch_exact): for name, p in model.named_parameters(): if "swav_prototypes" in name: p.grad = None - return model + + iters_noupdate = 0 + if cfg.MODEL.MODEL_NAME == "ContrastiveModel" and cfg.CONTRASTIVE.TYPE == "moco": + assert cfg.CONTRASTIVE.QUEUE_LEN % (cfg.TRAIN.BATCH_SIZE * cfg.NUM_SHARDS) == 0 + iters_noupdate = ( + cfg.CONTRASTIVE.QUEUE_LEN // cfg.TRAIN.BATCH_SIZE // cfg.NUM_SHARDS + ) + + if cur_iter < iters_noupdate and epoch_exact < 1: # for e.g. MoCo + logger.info("Not updating parameters {}/{}".format(cur_iter, iters_noupdate)) + update_param = False + else: + update_param = True + + return model, update_param + + +def contrastive_forward(model, cfg, inputs, index, time, epoch_exact, scaler): + if cfg.CONTRASTIVE.SEQUENTIAL: + perform_backward = False + mdl = getattr(model, "module", model) + keys = ( + mdl.compute_key_feat( + inputs, + compute_predictor_keys=False, + batched_inference=True if len(inputs) < 2 else False, + ) + if cfg.CONTRASTIVE.TYPE == "moco" or cfg.CONTRASTIVE.TYPE == "byol" + else [None] * len(inputs) + ) + for k, vid in enumerate(inputs): + other_keys = keys[:k] + keys[k + 1 :] + time_cur = torch.cat( + [ + time[:, k : k + 1, :], + time[:, :k, :], + time[:, k + 1 :, :], + ], + 1, + ) # q, kpre, kpost + vids = [vid] + if cfg.CONTRASTIVE.TYPE == "swav" or cfg.CONTRASTIVE.TYPE == "simclr": + if k < len(inputs) - 1: + vids = inputs[k : k + 2] + else: + break + lgt_k, loss_k = model(vids, index, time_cur, epoch_exact, keys=other_keys) + scaler.scale(loss_k).backward() + if k == 0: + preds, partial_loss = lgt_k, loss_k.detach() + else: + preds = torch.cat([preds, lgt_k], dim=0) + partial_loss += loss_k.detach() + partial_loss /= len(inputs) * 2.0 # to have same loss as symm model + if cfg.CONTRASTIVE.TYPE == "moco": + mdl._dequeue_and_enqueue(keys) + else: + perform_backward = True + preds, partial_loss = model(inputs, index, time, epoch_exact, keys=None) + return model, preds, partial_loss, perform_backward diff --git a/slowfast/models/head_helper.py b/slowfast/models/head_helper.py index 937208a7b..12f6eaed4 100644 --- a/slowfast/models/head_helper.py +++ b/slowfast/models/head_helper.py @@ -3,14 +3,18 @@ """ResNe(X)t Head helper.""" +from functools import partial + +import slowfast.utils.logging as logging import torch import torch.nn as nn from detectron2.layers import ROIAlign - +from slowfast.models.attention import MultiScaleBlock from slowfast.models.batchnorm_helper import ( NaiveSyncBatchNorm1d as NaiveSyncBatchNorm1d, ) -from slowfast.models.nonlocal_helper import Nonlocal + +logger = logging.get_logger(__name__) class ResNetRoIHead(nn.Module): @@ -71,16 +75,14 @@ def __init__( performance if ROIAlign is used together with conv layers. """ super(ResNetRoIHead, self).__init__() - assert ( - len({len(pool_size), len(dim_in)}) == 1 - ), "pathway dimensions are not consistent." + assert len({len(pool_size), len(dim_in)}) == 1, ( + "pathway dimensions are not consistent." + ) self.num_pathways = len(pool_size) self.detach_final_fc = detach_final_fc for pathway in range(self.num_pathways): - temporal_pool = nn.AvgPool3d( - [pool_size[pathway][0], 1, 1], stride=1 - ) + temporal_pool = nn.AvgPool3d([pool_size[pathway][0], 1, 1], stride=1) self.add_module("s{}_tpool".format(pathway), temporal_pool) roi_align = ROIAlign( @@ -107,14 +109,13 @@ def __init__( self.act = nn.Sigmoid() else: raise NotImplementedError( - "{} is not supported as an activation" - "function.".format(act_func) + "{} is not supported as an activationfunction.".format(act_func) ) def forward(self, inputs, bboxes): - assert ( - len(inputs) == self.num_pathways - ), "Input tensor does not contain {} pathway".format(self.num_pathways) + assert len(inputs) == self.num_pathways, ( + "Input tensor does not contain {} pathway".format(self.num_pathways) + ) pool_out = [] for pathway in range(self.num_pathways): t_pool = getattr(self, "s{}_tpool".format(pathway)) @@ -236,9 +237,9 @@ def __init__( cfg (struct): The config for the current experiment. """ super(ResNetBasicHead, self).__init__() - assert ( - len({len(pool_size), len(dim_in)}) == 1 - ), "pathway dimensions are not consistent." + assert len({len(pool_size), len(dim_in)}) == 1, ( + "pathway dimensions are not consistent." + ) self.num_pathways = len(pool_size) self.detach_final_fc = detach_final_fc self.cfg = cfg @@ -266,12 +267,10 @@ def __init__( cfg.CONTRASTIVE.MLP_DIM, cfg.CONTRASTIVE.NUM_MLP_LAYERS, bn_on=cfg.CONTRASTIVE.BN_MLP, - bn_sync_num=cfg.BN.NUM_SYNC_DEVICES - if cfg.CONTRASTIVE.BN_SYNC_MLP - else 1, - global_sync=( - cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC + bn_sync_num=( + cfg.BN.NUM_SYNC_DEVICES if cfg.CONTRASTIVE.BN_SYNC_MLP else 1 ), + global_sync=(cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC), ) # Softmax for evaluation and testing. @@ -283,13 +282,12 @@ def __init__( self.act = None else: raise NotImplementedError( - "{} is not supported as an activation" - "function.".format(act_func) + "{} is not supported as an activationfunction.".format(act_func) ) if cfg.CONTRASTIVE.PREDICTOR_DEPTHS: d_in = num_classes - for i, n_layers in enumerate(cfg.CONTRASTIVE.PREDICTOR_DEPTHS): + for n_layers in cfg.CONTRASTIVE.PREDICTOR_DEPTHS: local_mlp = MLPHead( d_in, num_classes, @@ -297,19 +295,17 @@ def __init__( n_layers, bn_on=cfg.CONTRASTIVE.BN_MLP, flatten=False, - bn_sync_num=cfg.BN.NUM_SYNC_DEVICES - if cfg.CONTRASTIVE.BN_SYNC_MLP - else 1, - global_sync=( - cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC + bn_sync_num=( + cfg.BN.NUM_SYNC_DEVICES if cfg.CONTRASTIVE.BN_SYNC_MLP else 1 ), + global_sync=(cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC), ) self.predictors.append(local_mlp) def forward(self, inputs): - assert ( - len(inputs) == self.num_pathways - ), "Input tensor does not contain {} pathway".format(self.num_pathways) + assert len(inputs) == self.num_pathways, ( + "Input tensor does not contain {} pathway".format(self.num_pathways) + ) pool_out = [] for pathway in range(self.num_pathways): m = getattr(self, "pathway{}_avgpool".format(pathway)) @@ -414,7 +410,6 @@ def __init__( self._construct_head(dim_in, dim_inner, dim_out, norm_module) def _construct_head(self, dim_in, dim_inner, dim_out, norm_module): - self.conv_5 = nn.Conv3d( dim_in, dim_inner, @@ -460,8 +455,7 @@ def _construct_head(self, dim_in, dim_inner, dim_out, norm_module): self.act = nn.Sigmoid() else: raise NotImplementedError( - "{} is not supported as an activation" - "function.".format(self.act_func) + "{} is not supported as an activationfunction.".format(self.act_func) ) def forward(self, inputs): @@ -531,12 +525,10 @@ def __init__( cfg.CONTRASTIVE.MLP_DIM, cfg.CONTRASTIVE.NUM_MLP_LAYERS, bn_on=cfg.CONTRASTIVE.BN_MLP, - bn_sync_num=cfg.BN.NUM_SYNC_DEVICES - if cfg.CONTRASTIVE.BN_SYNC_MLP - else 1, - global_sync=( - cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC + bn_sync_num=( + cfg.BN.NUM_SYNC_DEVICES if cfg.CONTRASTIVE.BN_SYNC_MLP else 1 ), + global_sync=(cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC), ) self.detach_final_fc = cfg.MODEL.DETACH_FINAL_FC @@ -549,8 +541,7 @@ def __init__( self.act = None else: raise NotImplementedError( - "{} is not supported as an activation" - "function.".format(act_func) + "{} is not supported as an activationfunction.".format(act_func) ) def forward(self, x): @@ -570,3 +561,112 @@ def forward(self, x): x = x.view(x.shape[0], -1) return x + + +class MSSeparateHead(nn.Module): + """ + Perform linear projection or Transformer-based decoder (optionally MultiScale) + for mask prediction models. + Args: + blocks (MultiScaleBlock): the encoder blocks to provide input dimensions of the head. + num_classes (int): the dimension of the prediction target (eg. HOG or pixels). + feat_sz (list): the spatiotemporal sizes of the input features. + """ + + def __init__( + self, + blocks, + cfg, + num_classes, + feat_sz, + ): + super(MSSeparateHead, self).__init__() + head_type = cfg.MASK.HEAD_TYPE.split("_") + assert head_type[0] == "separate" + if len(head_type) > 1: + transform_type = head_type[1] + assert transform_type in ["xformer"] + else: + transform_type = None + + depth_list = cfg.MASK.PRETRAIN_DEPTH + mlp_ratio = cfg.MVIT.MLP_RATIO + qkv_bias = cfg.MVIT.QKV_BIAS + drop_rate = cfg.MVIT.DROPOUT_RATE + kernel_kv = cfg.MASK.DEC_KV_KERNEL + stride_kv = cfg.MASK.DEC_KV_STRIDE + mode = cfg.MVIT.MODE + self.cls_embed_on = cfg.MVIT.CLS_EMBED_ON + pool_first = cfg.MVIT.POOL_FIRST + if cfg.MVIT.NORM == "layernorm": + norm_layer = partial(nn.LayerNorm, eps=1e-6) + else: + raise NotImplementedError("Only supports layernorm.") + + self.transforms = nn.ModuleList() + self.projections = nn.ModuleList() + for depth, num_class, feature_size in zip(depth_list, num_classes, feat_sz): + head_dim = ( + cfg.MASK.DECODER_EMBED_DIM if cfg.MASK.MAE_ON else blocks[depth].dim_out + ) + op = [] + if transform_type == "xformer": + assert cfg.MASK.DECODER_DEPTH > 0 + for _ in range(cfg.MASK.DECODER_DEPTH): + dim_out = cfg.MASK.DECODER_EMBED_DIM + op.append( + MultiScaleBlock( + dim=head_dim, + dim_out=dim_out, + input_size=feature_size, + num_heads=dim_out // 64, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop_rate=drop_rate, + drop_path=0.0, + norm_layer=norm_layer, + kernel_q=[], + kernel_kv=kernel_kv, + stride_q=[], + stride_kv=stride_kv, + mode=mode, + has_cls_embed=self.cls_embed_on, + pool_first=pool_first, + ) + ) + head_dim = dim_out + + op.append(norm_layer(head_dim)) + self.transforms.append(nn.Sequential(*op)) + self.projections.append(nn.Linear(head_dim, num_class, bias=True)) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + # FYI: MAE uses xavier_uniform following official JAX ViT: + # torch.nn.init.xavier_uniform_(m.weight) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def forward(self, block_outputs, output_masks, return_all, thw): + model_outputs = [] + for idx, x in enumerate(block_outputs): + for blk in self.transforms[idx]: + if isinstance(blk, MultiScaleBlock): + x, thw = blk(x, thw) + else: + x = blk(x) + + if self.cls_embed_on: + x = x[:, 1:] + if not return_all: + mask = output_masks[idx] + x = x[mask] + x = self.projections[idx](x) + model_outputs.append(x) + return model_outputs diff --git a/slowfast/models/losses.py b/slowfast/models/losses.py index 96e39b5f4..6f419180c 100644 --- a/slowfast/models/losses.py +++ b/slowfast/models/losses.py @@ -7,10 +7,7 @@ import torch import torch.nn as nn - -from pytorchvideo.losses.soft_target_cross_entropy import ( - SoftTargetCrossEntropyLoss, -) +from pytorchvideo.losses.soft_target_cross_entropy import SoftTargetCrossEntropyLoss class ContrastiveLoss(nn.Module): @@ -20,18 +17,55 @@ def __init__(self, reduction="mean"): def forward(self, inputs, dummy_labels=None): targets = torch.zeros(inputs.shape[0], dtype=torch.long).cuda() - loss = nn.CrossEntropyLoss(reduction=self.reduction).cuda()( - inputs, targets - ) + loss = nn.CrossEntropyLoss(reduction=self.reduction).cuda()(inputs, targets) return loss +class MultipleMSELoss(nn.Module): + """ + Compute multiple mse losses and return their average. + """ + + def __init__(self, reduction="mean"): + """ + Args: + reduction (str): specifies reduction to apply to the output. It can be + "mean" (default) or "none". + """ + super(MultipleMSELoss, self).__init__() + self.mse_func = nn.MSELoss(reduction=reduction) + + def forward(self, x, y): + loss_sum = 0.0 + multi_loss = [] + for xt, yt in zip(x, y): + if isinstance(yt, (tuple,)): + if len(yt) == 2: + yt, wt = yt + lt = "mse" + elif len(yt) == 3: + yt, wt, lt = yt + else: + raise NotImplementedError + else: + wt, lt = 1.0, "mse" + if lt == "mse": + loss = self.mse_func(xt, yt) + else: + raise NotImplementedError + loss_sum += loss * wt + multi_loss.append(loss) + return loss_sum, multi_loss + + _LOSSES = { "cross_entropy": nn.CrossEntropyLoss, "bce": nn.BCELoss, "bce_logit": nn.BCEWithLogitsLoss, "soft_cross_entropy": partial(SoftTargetCrossEntropyLoss, normalize_targets=False), "contrastive_loss": ContrastiveLoss, + "mse": nn.MSELoss, + "multi_mse": MultipleMSELoss, } diff --git a/slowfast/models/masked.py b/slowfast/models/masked.py new file mode 100644 index 000000000..839f1ff30 --- /dev/null +++ b/slowfast/models/masked.py @@ -0,0 +1,623 @@ +#!/usr/bin/env python3 +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. + +import math +from functools import partial + +import slowfast.utils.logging as logging +import slowfast.utils.misc as misc +import torch +import torch.nn as nn +import torch.nn.functional as F +from slowfast.models import head_helper +from slowfast.models.utils import calc_mvit_feature_geometry +from slowfast.models.video_model_builder import MViT +from torch.nn.init import trunc_normal_ + +from . import head_helper, operators, resnet_helper, stem_helper # noqa +from .build import MODEL_REGISTRY + +logger = logging.get_logger(__name__) + + +@MODEL_REGISTRY.register() +class MaskMViT(MViT): + def __init__(self, cfg): + super().__init__(cfg) + self.pretrain_depth = cfg.MASK.PRETRAIN_DEPTH + if self.pretrain_depth[-1] + 1 < cfg.MVIT.DEPTH: + del self.blocks[self.pretrain_depth[-1] + 1 :] + del self.norm + del self.head + self.feat_size, self.feat_stride = calc_mvit_feature_geometry(cfg) + + self.head_type = cfg.MASK.HEAD_TYPE.split("_") + feat_sz = [self.feat_size[depth] for depth in self.pretrain_depth] + if self.head_type[0] == "separate": + if not cfg.MASK.PRED_HOG: + pred_t_sz = ( + 1 if self.cfg.MASK.TIME_STRIDE_LOSS else self.patch_stride[0] + ) + num_classes = [ + pred_t_sz * (self.feat_stride[depth][-1] ** 2) * 3 + for depth in self.pretrain_depth + ] + self.pred_head = head_helper.MSSeparateHead( + self.blocks, cfg, num_classes, feat_sz + ) + else: + self.hogs = nn.ModuleList() + self.nbins = 9 + self.cell_sz = 8 + self.hogs.append( + operators.HOGLayerC( + nbins=self.nbins, + pool=self.cell_sz, + ) + ) + self.ncells = [ + (self.feat_stride[depth][-1] // self.cell_sz) ** 2 + for depth in self.pretrain_depth + ] + pred_hog_classes = [self.nbins * ncell for ncell in self.ncells] + pred_hog_classes = [ + pred_hog_class * 3 # 3 color channels + for pred_hog_class in pred_hog_classes + ] + self.pred_head = head_helper.MSSeparateHead( + self.blocks, cfg, pred_hog_classes, feat_sz + ) + self.hog_loss = "mse" + else: + raise NotImplementedError + + embed_dim = cfg.MVIT.EMBED_DIM + decoder_embed_dim = cfg.MASK.DECODER_EMBED_DIM + self.sep_pos_embed_decoder = cfg.MASK.DECODER_SEP_POS_EMBED + self.counter = 0 + if cfg.MASK.MAE_ON: + # ---------------------------------------------------------------- + # MAE decoder specifics + norm_layer = partial(nn.LayerNorm, eps=1e-6) + dim_in = self.blocks[-1].dim_out + self.norm = norm_layer(dim_in) + self.decoder_embed = nn.Linear(dim_in, decoder_embed_dim, bias=True) + num_patches = math.prod(self.patch_dims) + if self.use_abs_pos: + if self.sep_pos_embed_decoder: + self.dec_pos_embed_spatial = nn.Parameter( + torch.zeros( + 1, + self.patch_dims[1] * self.patch_dims[2], + decoder_embed_dim, + ) + ) + self.dec_pos_embed_temporal = nn.Parameter( + torch.zeros(1, self.patch_dims[0], decoder_embed_dim) + ) + if self.cls_embed_on: + self.dec_pos_embed_class = nn.Parameter( + torch.zeros(1, 1, decoder_embed_dim) + ) + else: + self.decoder_pos_embed = nn.Parameter( + torch.zeros( + 1, + num_patches + 1 if self.cls_embed_on else num_patches, + decoder_embed_dim, + ) + ) + self.mask_token = nn.Parameter( + torch.zeros(1, 1, decoder_embed_dim if cfg.MASK.MAE_ON else embed_dim) + ) + trunc_normal_(self.mask_token, std=0.02) + if self.use_abs_pos and cfg.MASK.MAE_ON: + if self.sep_pos_embed_decoder: + trunc_normal_(self.dec_pos_embed_spatial, std=0.02) + trunc_normal_(self.dec_pos_embed_temporal, std=0.02) + if self.cls_embed_on: + trunc_normal_(self.dec_pos_embed_class, std=0.02) + else: + trunc_normal_(self.decoder_pos_embed, std=0.02) + + if cfg.MASK.SCALE_INIT_BY_DEPTH: + self.fix_init_weight() + + self.pred_pixel_wt = 0.0 if cfg.MASK.PRED_HOG else 1.0 + self.pred_hog_wt = 1.0 if cfg.MASK.PRED_HOG else 0.0 + + @torch.jit.ignore + def no_weight_decay(self): + names = [] + if self.cfg.MVIT.ZERO_DECAY_POS_CLS: + if self.use_abs_pos: + if self.sep_pos_embed_decoder: + names.extend( + [ + "dec_pos_embed_spatial", + "dec_pos_embed_temporal", + "dec_pos_embed_class", + ] + ) + else: + names.extend(["pos_embed_decoder"]) + if self.cls_embed_on: + names.append("cls_token") + + return names + + def fix_init_weight(self): + def rescale(param, layer_id): + param.div_(math.sqrt(2.0 * layer_id)) + + for layer_id, layer in enumerate(self.blocks): + rescale(layer.attn.proj.weight.data, layer_id + 1) + rescale(layer.mlp.fc2.weight.data, layer_id + 1) + for trans in self.pred_head.transforms: + for layer_id, layer in enumerate(trans): + if hasattr(layer, "attn"): + rescale( + layer.attn.proj.weight.data, + layer_id + 1 + len(self.blocks), + ) # or + len(self.blocks) + rescale(layer.mlp.fc2.weight.data, layer_id + 1) + + def _get_multiscale_mask(self, mask): + if self.use_2d_patch: + mask = mask.unsqueeze(0) + output_masks = [] + for depth in self.pretrain_depth: + size = self.feat_size[depth][-1] + output_mask = F.interpolate(mask, size=size) + if self.use_2d_patch: + output_mask = output_mask[0] + output_mask = output_mask.flatten(1).to(torch.bool) + output_masks.append(output_mask) + return output_masks + + def _patchify(self, imgs, p=16, time_stride_loss=True): + N, _, T, H, W = imgs.shape + u = 1 if time_stride_loss else self.patch_stride[0] + assert H == W and H % p == 0 and T % u == 0 + h = w = H // p + t = T // u + x = imgs.reshape(shape=(N, 3, t, u, h, p, w, p)) + x = torch.einsum("nctuhpwq->nthwupqc", x) + x = x.reshape(shape=(N, t * h * w, u * p**2 * 3)) + self.patch_info = (N, T, H, W, p, u, t, h, w) + return x + + def _unpatchify(self, x): + N, T, H, W, p, u, t, h, w = self.patch_info + x = x.reshape(shape=(N, t, h, w, u, p, p, 3)) + x = torch.einsum("nthwupqc->nctuhpwq", x) + imgs = x.reshape(shape=(N, 3, T, H, W)) + return imgs + + def _get_pixel_label_2d(self, input_img, output_masks, norm=True): + input_img = input_img.permute(0, 2, 3, 1) + labels = [] + for depth, output_mask in zip(self.pretrain_depth, output_masks): + size = self.feat_stride[depth][-1] + label = input_img.unfold(1, size, size).unfold(2, size, size) + label = label.flatten(1, 2).flatten(2) + label = label[output_mask] + if norm: + mean = label.mean(dim=-1, keepdim=True) + var = label.var(dim=-1, keepdim=True) + label = (label - mean) / (var + 1.0e-6) ** 0.5 + labels.append(label) + return labels + + def _get_pixel_label_3d( + self, input_frames, output_masks, time_stride_loss=True, norm=True + ): + if time_stride_loss: + input_frames = input_frames[:, :, :: self.cfg.MVIT.PATCH_STRIDE[0], :, :] + imgs = input_frames + input_frames = input_frames.permute(0, 2, 3, 4, 1) + labels = [] + for depth, output_mask in zip(self.pretrain_depth, output_masks): + size = self.feat_stride[depth][-1] + label = self._patchify(imgs, p=size, time_stride_loss=time_stride_loss) + label = label[output_mask] + + if norm: # self.norm_pix_loss: + mean = label.mean(dim=-1, keepdim=True) + var = label.var(dim=-1, keepdim=True) + label = (label - mean) / (var + 1.0e-6) ** 0.5 + labels.append((label, self.pred_pixel_wt / len(self.pretrain_depth))) + return labels + + def _get_hog_label_2d(self, input_frames, output_masks): + # input_frames, B C H W + labels = [] + for depth, output_mask in zip(self.pretrain_depth, output_masks): + feat_size = self.feat_size[depth][-1] + hog_list = [] + for hog in self.hogs: + tmp_hog = hog(input_frames).flatten(1, 2) # return B C H W + unfold_size = tmp_hog.shape[-1] // feat_size + tmp_hog = ( + tmp_hog.permute(0, 2, 3, 1) + .unfold(1, unfold_size, unfold_size) + .unfold(2, unfold_size, unfold_size) + .flatten(1, 2) + .flatten(2) + ) + tmp_hog = tmp_hog[output_mask] + hog_list.append(tmp_hog) + all_tlabel = torch.cat(hog_list, -1) + labels.append((all_tlabel, self.pred_hog_wt, self.hog_loss)) + return labels + + def _get_hog_label_3d(self, input_frames, output_masks): + input_frames = input_frames[ + :, :, :: self.cfg.MVIT.PATCH_STRIDE[0], :, : + ] # B C T H W + input_frames = input_frames.transpose(1, 2) # B T C H W + B, T = input_frames.shape[:2] + input_frames = input_frames.flatten(0, 1) # BT C H W + labels = [] + for depth, output_mask in zip(self.pretrain_depth, output_masks): + feat_size = self.feat_size[depth][-1] + hog_list = [] + for hog in self.hogs: + tmp_hog = hog(input_frames).flatten(1, 2) # BT C H W + unfold_size = tmp_hog.shape[-1] // feat_size + tmp_hog = ( + tmp_hog.permute(0, 2, 3, 1) + .unfold(1, unfold_size, unfold_size) + .unfold(2, unfold_size, unfold_size) + ) # BT h w C wh ww + tmp_hog = tmp_hog.flatten(3).view( + B, T, feat_size, feat_size, -1 + ) # B T h w C (3 nbins h w) + tmp_hog = tmp_hog.flatten(1, 3) # B N C + tmp_hog = tmp_hog[output_mask] + hog_list.append(tmp_hog) + all_tlabel = torch.cat(hog_list, -1) + labels.append((all_tlabel, self.pred_hog_wt, self.hog_loss)) + return labels + + def _mae_random_masking(self, x, mask_ratio, mask_in=None): + """ + Perform per-sample random masking by per-sample shuffling. + Per-sample shuffling is done by argsort random noise. + x: [N, L, D], sequence + """ + N, L, D = x.shape # batch, length, dim + if mask_in is None: + if self.cfg.AUG.MASK_TUBE: + noise = ( + torch.rand(N, 1, self.H * self.W, device=x.device) + .repeat([1, self.T, 1]) + .reshape(N, L) + ) # noise in [0, 1] + else: + noise = torch.rand(N, L, device=x.device) # noise in [0, 1] + else: + noise = mask_in.flatten(1) + mask_ratio = sum(noise.flatten()) / noise.numel() # alrdy masked + len_keep = int(L * (1 - mask_ratio)) + assert len_keep > 1 + # sort noise for each sample + ids_shuffle = torch.argsort( + noise, dim=1 + ) # ascend: small is keep, large is remove + ids_restore = torch.argsort(ids_shuffle, dim=1) + # keep the first subset + ids_keep = ids_shuffle[:, :len_keep] + x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) + # generate the binary mask: 0 is keep, 1 is remove + mask = torch.ones([N, L], device=x.device) + mask[:, :len_keep] = 0 + # unshuffle to get the binary mask + mask = torch.gather(mask, dim=1, index=ids_restore) + return x_masked, mask, ids_restore, ids_keep + + def _mae_forward_encoder(self, x, mask_ratio, mask=None): + x, bcthw = self.patch_embed(x, keep_spatial=False) + bcthw = list(bcthw) + if len(bcthw) == 4: # Fix bcthw in case of 4D tensor + bcthw.insert(2, torch.tensor(self.T)) + T, H, W = bcthw[-3], bcthw[-2], bcthw[-1] + assert len(bcthw) == 5 and (T, H, W) == (self.T, self.H, self.W), bcthw + s = 1 if self.cls_embed_on else 0 + B, N, C = x.shape + + if self.use_fixed_sincos_pos: + x += self.pos_embed[:, s:, :] # 0: no cls token + + if self.cfg.MASK.PER_FRAME_MASKING: + x = x.reshape([B * T, H * W, C]) + x, mask, ids_restore, ids_keep = self._mae_random_masking( + x, mask_ratio, None if self.cfg.MASK.MAE_RND_MASK else mask + ) + if self.cfg.MASK.PER_FRAME_MASKING: + x = x.view([B, -1, C]) + + if self.cls_embed_on: + # append cls token + cls_token = self.cls_token # + if self.use_fixed_sincos_pos: + cls_token = cls_token + self.pos_embed[:, :s, :] + cls_tokens = cls_token.expand(x.shape[0], -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + + if self.use_abs_pos and not self.use_fixed_sincos_pos: + if self.sep_pos_embed: + pos_embed = self.pos_embed_spatial.repeat( + 1, self.patch_dims[0], 1 + ) + torch.repeat_interleave( + self.pos_embed_temporal, + self.patch_dims[1] * self.patch_dims[2], + dim=1, + ) + pos_embed = pos_embed.expand(x.shape[0], -1, -1) + pos_embed = torch.gather( + pos_embed, + dim=1, + index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_embed.shape[2]), + ) + if self.cls_embed_on: + pos_embed = torch.cat( + [ + self.pos_embed_class.expand(pos_embed.shape[0], -1, -1), + pos_embed, + ], + 1, + ) + x += pos_embed + else: + pos_embed = self.pos_embed.expand(x.shape[0], -1, -1) + pos_embed_sampled = torch.gather( + pos_embed[:, s:, :], + dim=1, + index=ids_keep.unsqueeze(-1).repeat(1, 1, self.pos_embed.shape[2]), + ) + if self.cls_embed_on: + pos_embed_sampled = torch.cat( + [pos_embed[:, :s, :], pos_embed_sampled], 1 + ) + x += pos_embed_sampled + + # apply Transformer blocks + B, N, C = x.shape + thw = [T, H, W] + for _, blk in enumerate(self.blocks): + x, thw = blk(x, thw) + x = self.norm(x) + + return x, mask, ids_restore, thw + + def _mae_forward_decoder(self, x, ids_restore, mask, thw): + # embed tokens + x = self.decoder_embed(x) + T, H, W = self.T, self.H, self.W + B, N, C = x.shape + + s = 1 if self.cls_embed_on else 0 + + # append mask tokens to sequence + mask_tokens = self.mask_token.repeat( + B, T * H * W + s - x.shape[1], 1 + ) # + s: no cls token + x_ = torch.cat([x[:, s:, :], mask_tokens], dim=1) # no cls token + if self.cfg.MASK.PER_FRAME_MASKING: + x_ = x_.view([B * T, H * W, C]) + x_ = torch.gather( + x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x_.shape[2]) + ) # unshuffle + if self.cfg.MASK.PER_FRAME_MASKING: + x_ = x_.view([B, T * H * W, C]) + x = torch.cat([x[:, :s, :], x_], dim=1) # append cls token + + if self.sep_pos_embed_decoder: + pos_embed = self.dec_pos_embed_spatial.repeat( + 1, self.patch_dims[0], 1 + ) + torch.repeat_interleave( + self.dec_pos_embed_temporal, + self.patch_dims[1] * self.patch_dims[2], + dim=1, + ) + pos_embed = pos_embed.expand(x.shape[0], -1, -1) + if self.cls_embed_on: + pos_embed = torch.cat( + [ + self.dec_pos_embed_class.expand(pos_embed.shape[0], -1, -1), + pos_embed, + ], + 1, + ) + x += pos_embed + else: + # add pos embed + x = x + self.decoder_pos_embed + + pixel_outputs = self.pred_head( + [x], + [mask.to(torch.bool)], + return_all=self.cfg.VIS_MASK.ENABLE, + thw=thw, + ) + + return pixel_outputs + + def _mae_forward(self, imgs, mask_ratio=0.75, mask=None): + latent, mask, ids_restore, thw = self._mae_forward_encoder( + imgs, mask_ratio, mask + ) + pred = self._mae_forward_decoder(latent, ids_restore, mask, thw) + labels = [] + if self.pred_pixel_wt: + if self.use_2d_patch: + labels += self._get_pixel_label_2d( + imgs.detach(), + [mask.to(torch.bool)], + norm=self.cfg.MASK.NORM_PRED_PIXEL, + ) + else: + labels += self._get_pixel_label_3d( + imgs.detach(), + [mask.to(torch.bool)], + time_stride_loss=self.cfg.MASK.TIME_STRIDE_LOSS, + norm=self.cfg.MASK.NORM_PRED_PIXEL, + ) + if self.pred_hog_wt: + if self.use_2d_patch: + labels += self._get_hog_label_2d(imgs.detach(), [mask.to(torch.bool)]) + else: + labels += self._get_hog_label_3d(imgs.detach(), [mask.to(torch.bool)]) + + self.counter += 1 + if self.cfg.VIS_MASK.ENABLE: + return self._mae_visualize(imgs, pred, mask) + return pred, labels + + def _mae_visualize(self, imgs, pred, mask): + N, T, H, W, p, u, t, h, w = self.patch_info + pred = pred[0] + if self.cfg.MASK.TIME_STRIDE_LOSS: + im_viz = imgs[:, :, :: self.cfg.MVIT.PATCH_STRIDE[0], :, :] + else: + im_viz = imgs + reconstruct = self._unpatchify( + pred * mask.reshape(N, t * h * w, 1) + + self._patchify(im_viz, time_stride_loss=self.cfg.MASK.TIME_STRIDE_LOSS) + * (1 - mask.reshape(N, t * h * w, 1)) + ) + masked = self._unpatchify( + self._patchify(im_viz, time_stride_loss=self.cfg.MASK.TIME_STRIDE_LOSS) + * (1 - mask.reshape(N, t * h * w, 1)) + ) + + comparison = torch.stack( + [im_viz, masked, reconstruct], + dim=1, + ).permute([0, 1, 3, 2, 4, 5]) + pfx = self.cfg.TEST.CHECKPOINT_FILE_PATH + mr = self.cfg.AUG.MASK_RATIO + for i in range(comparison.shape[0]): + misc.plot_input_normed( + comparison[i].cpu(), + bboxes=(), + texts=(), + path=self.cfg.OUTPUT_DIR + + "/vis_mask/vid/{}vis_video_in_mask_out_mr{}/vis_{}_{}.mp4".format( + pfx[pfx.rfind("/") + 1 : -5], mr, self.counter, i + ), + folder_path=self.cfg.OUTPUT_DIR + + "/vis_mask/vid/{}vis_video_in_mask_out_mr{}".format( + pfx[pfx.rfind("/") + 1 : -5], mr + ), + make_grids=True, + output_video=True, + ) + return pred[0] + + def _maskfeat_forward(self, x, mask, return_all=False): + x_embed, x_shape = self.patch_embed(x) + if self.cfg.MASK.MAE_RND_MASK: + _, mask, ids_restore, ids_keep = self._mae_random_masking( + x_embed, self.cfg.AUG.MASK_RATIO, None + ) + output_masks = [mask.to(torch.bool)] + else: + # take masks and labels from loader + float_mask = mask.type_as(x) + output_masks = self._get_multiscale_mask(float_mask) + labels = [] + if self.pred_pixel_wt: + if self.use_2d_patch: + labels += self._get_pixel_label_2d( + x.detach(), output_masks, norm=self.cfg.MASK.NORM_PRED_PIXEL + ) + else: + labels += self._get_pixel_label_3d( + x.detach(), output_masks, norm=self.cfg.MASK.NORM_PRED_PIXEL + ) + if self.pred_hog_wt: + if self.use_2d_patch: + labels += self._get_hog_label_2d(x.detach(), output_masks) + else: + labels += self._get_hog_label_3d(x.detach(), output_masks) + + x = x_embed + T, H, W = self.T, self.H, self.W + B, N, C = x.shape + + # switch input tokens by mask_token + mask_tokens = self.mask_token.expand(B, N, -1) + if self.cfg.MASK.MAE_RND_MASK: + float_mask = mask.unsqueeze(-1) + else: + if self.use_2d_patch: + float_mask = F.interpolate(float_mask.unsqueeze(0), size=(H, W))[0] + else: + float_mask = F.interpolate(float_mask, size=(H, W)) + float_mask = float_mask.flatten(1).unsqueeze(-1) + x = x * (1 - float_mask) + mask_tokens * float_mask + + if self.cls_embed_on: + cls_tokens = self.cls_token.expand(B, -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + + if self.use_abs_pos: + if self.sep_pos_embed: + pos_embed = self.pos_embed_spatial.repeat( + 1, self.patch_dims[0], 1 + ) + torch.repeat_interleave( + self.pos_embed_temporal, + self.patch_dims[1] * self.patch_dims[2], + dim=1, + ) + if self.cls_embed_on: + pos_embed = torch.cat([self.pos_embed_class, pos_embed], 1) + x = x + pos_embed + else: + x = x + self.pos_embed + + if self.drop_rate: + x = self.pos_drop(x) + + if self.norm_stem: + x = self.norm_stem(x) + + thw = [T, H, W] + block_outputs = [] + for idx, blk in enumerate(self.blocks): + x, thw = blk(x, thw) + if idx in self.pretrain_depth: + block_outputs.append(x) + + model_outputs = [] + if self.pred_pixel_wt: + pixel_outputs = self.pred_head( + block_outputs, + output_masks, + return_all=return_all, + thw=thw, + ) + model_outputs += pixel_outputs + if self.pred_hog_wt: + hog_outputs = self.pred_head( + block_outputs, + output_masks, + return_all=return_all, + thw=thw, + ) + model_outputs += hog_outputs + + return model_outputs, labels + + def forward(self, x, return_all=False): + if len(x) > 1: + x, meta, mask = x + else: + x, mask = x[0], None + + if self.cfg.MASK.MAE_ON: + return self._mae_forward(x, mask_ratio=self.cfg.AUG.MASK_RATIO, mask=mask) + else: + return self._maskfeat_forward(x, mask, return_all) diff --git a/slowfast/models/nonlocal_helper.py b/slowfast/models/nonlocal_helper.py index dc699ccee..b84765140 100644 --- a/slowfast/models/nonlocal_helper.py +++ b/slowfast/models/nonlocal_helper.py @@ -54,9 +54,7 @@ def __init__( self.pool_size = pool_size self.instantiation = instantiation self.use_pool = ( - False - if pool_size is None - else any((size > 1 for size in pool_size)) + False if pool_size is None else any((size > 1 for size in pool_size)) ) self.norm_eps = norm_eps self.norm_momentum = norm_momentum @@ -133,9 +131,7 @@ def forward(self, x): spatial_temporal_dim = theta_phi.shape[2] theta_phi = theta_phi / spatial_temporal_dim else: - raise NotImplementedError( - "Unknown norm type {}".format(self.instantiation) - ) + raise NotImplementedError("Unknown norm type {}".format(self.instantiation)) # (N, TxHxW, TxHxW) * (N, C, TxHxW) => (N, C, TxHxW). theta_phi_g = torch.einsum("ntg,ncg->nct", (theta_phi, g)) diff --git a/slowfast/models/operators.py b/slowfast/models/operators.py index 1e4a0433e..6f0e4de45 100644 --- a/slowfast/models/operators.py +++ b/slowfast/models/operators.py @@ -3,9 +3,13 @@ """Custom operators.""" -import torch.nn as nn +import math +import torch +import torch.nn as nn +import torch.nn.functional as F from pytorchvideo.layers.swish import Swish +from slowfast.models.utils import get_gkern class SE(nn.Module): @@ -25,9 +29,7 @@ def _round_width(self, width, multiplier, min_width=8, divisor=8): width *= multiplier min_width = min_width or divisor - width_out = max( - min_width, int(width + divisor / 2) // divisor * divisor - ) + width_out = max(min_width, int(width + divisor / 2) // divisor * divisor) if width_out < 0.9 * width: width_out += divisor return int(width_out) @@ -55,3 +57,56 @@ def forward(self, x): for module in self.children(): x = module(x) return x_in * x + + +class HOGLayerC(nn.Module): + def __init__(self, nbins=9, pool=7, gaussian_window=0): + super(HOGLayerC, self).__init__() + self.nbins = nbins + self.pool = pool + self.pi = math.pi + weight_x = torch.FloatTensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]) + weight_x = weight_x.view(1, 1, 3, 3).repeat(3, 1, 1, 1) + weight_y = weight_x.transpose(2, 3) + self.register_buffer("weight_x", weight_x) + self.register_buffer("weight_y", weight_y) + + self.gaussian_window = gaussian_window + if gaussian_window: + gkern = get_gkern(gaussian_window, gaussian_window // 2) + self.register_buffer("gkern", gkern) + + @torch.no_grad() + def forward(self, x): + # input is RGB image with shape [B 3 H W] + x = F.pad(x, pad=(1, 1, 1, 1), mode="reflect") + gx_rgb = F.conv2d(x, self.weight_x, bias=None, stride=1, padding=0, groups=3) + gy_rgb = F.conv2d(x, self.weight_y, bias=None, stride=1, padding=0, groups=3) + norm_rgb = torch.stack([gx_rgb, gy_rgb], dim=-1).norm(dim=-1) + phase = torch.atan2(gx_rgb, gy_rgb) + phase = phase / self.pi * self.nbins # [-9, 9] + + b, c, h, w = norm_rgb.shape + out = torch.zeros((b, c, self.nbins, h, w), dtype=torch.float, device=x.device) + phase = phase.view(b, c, 1, h, w) + norm_rgb = norm_rgb.view(b, c, 1, h, w) + if self.gaussian_window: + if h != self.gaussian_window: + assert h % self.gaussian_window == 0, "h {} gw {}".format( + h, self.gaussian_window + ) + repeat_rate = h // self.gaussian_window + temp_gkern = self.gkern.repeat([repeat_rate, repeat_rate]) + else: + temp_gkern = self.gkern + norm_rgb *= temp_gkern + + out.scatter_add_(2, phase.floor().long() % self.nbins, norm_rgb) + + out = out.unfold(3, self.pool, self.pool) + out = out.unfold(4, self.pool, self.pool) + out = out.sum(dim=[-1, -2]) + + out = torch.nn.functional.normalize(out, p=2, dim=2) + + return out # B 3 nbins H W diff --git a/slowfast/models/optimizer.py b/slowfast/models/optimizer.py index 0ca5c4721..176283263 100644 --- a/slowfast/models/optimizer.py +++ b/slowfast/models/optimizer.py @@ -3,9 +3,8 @@ """Optimizer.""" -import torch - import slowfast.utils.lr_policy as lr_policy +import torch def construct_optimizer(model, cfg): @@ -23,74 +22,85 @@ def construct_optimizer(model, cfg): cfg (config): configs of hyper-parameters of SGD or ADAM, includes base learning rate, momentum, weight_decay, dampening, and etc. """ - bn_parameters = [] - non_bn_parameters = [] - zero_parameters = [] - no_grad_parameters = [] - skip = {} - if cfg.NUM_GPUS > 1: - if hasattr(model.module, "no_weight_decay"): - skip = model.module.no_weight_decay() - skip = {"module." + v for v in skip} + if cfg.SOLVER.LAYER_DECAY > 0.0 and cfg.SOLVER.LAYER_DECAY < 1.0: + optim_params = get_param_groups(model, cfg) + elif cfg.SOLVER.LAYER_DECAY == 1.0: + bn_parameters = [] + non_bn_parameters = [] + zero_parameters = [] + no_grad_parameters = [] + skip = {} + + if cfg.NUM_GPUS > 1: + if hasattr(model.module, "no_weight_decay"): + skip = model.module.no_weight_decay() + else: + if hasattr(model, "no_weight_decay"): + skip = model.no_weight_decay() + + for name_m, m in model.named_modules(): + is_bn = isinstance(m, torch.nn.modules.batchnorm._NormBase) + for name_p, p in m.named_parameters(recurse=False): + name = "{}.{}".format(name_m, name_p).strip(".") + if not p.requires_grad: + no_grad_parameters.append(p) + elif is_bn: + bn_parameters.append(p) + elif any(k in name for k in skip): + zero_parameters.append(p) + elif cfg.SOLVER.ZERO_WD_1D_PARAM and ( + len(p.shape) == 1 or name.endswith(".bias") + ): + zero_parameters.append(p) + else: + non_bn_parameters.append(p) + + optim_params = [ + { + "params": bn_parameters, + "weight_decay": cfg.BN.WEIGHT_DECAY, + "layer_decay": 1.0, + "apply_LARS": False, + }, + { + "params": non_bn_parameters, + "weight_decay": cfg.SOLVER.WEIGHT_DECAY, + "layer_decay": 1.0, + "apply_LARS": cfg.SOLVER.LARS_ON, + }, + { + "params": zero_parameters, + "weight_decay": 0.0, + "layer_decay": 1.0, + "apply_LARS": cfg.SOLVER.LARS_ON, + }, + ] + optim_params = [x for x in optim_params if len(x["params"])] + + # Check all parameters will be passed into optimizer. + assert len(list(model.parameters())) == len(non_bn_parameters) + len( + bn_parameters + ) + len(zero_parameters) + len(no_grad_parameters), ( + "parameter size does not match: {} + {} + {} + {} != {}".format( + len(non_bn_parameters), + len(bn_parameters), + len(zero_parameters), + len(no_grad_parameters), + len(list(model.parameters())), + ) + ) + print( + "bn {}, non bn {}, zero {}, no grad {}".format( + len(bn_parameters), + len(non_bn_parameters), + len(zero_parameters), + len(no_grad_parameters), + ) + ) else: - if hasattr(model, "no_weight_decay"): - skip = model.no_weight_decay() - - for name, m in model.named_modules(): - is_bn = isinstance(m, torch.nn.modules.batchnorm._NormBase) - for p in m.parameters(recurse=False): - if not p.requires_grad: - no_grad_parameters.append(p) - elif is_bn: - bn_parameters.append(p) - elif name in skip: - zero_parameters.append(p) - elif cfg.SOLVER.ZERO_WD_1D_PARAM and ( - len(p.shape) == 1 or name.endswith(".bias") - ): - zero_parameters.append(p) - else: - non_bn_parameters.append(p) - - optim_params = [ - { - "params": bn_parameters, - "weight_decay": cfg.BN.WEIGHT_DECAY, - "apply_LARS": False, - }, - { - "params": non_bn_parameters, - "weight_decay": cfg.SOLVER.WEIGHT_DECAY, - "apply_LARS": cfg.SOLVER.LARS_ON, - }, - { - "params": zero_parameters, - "weight_decay": 0.0, - "apply_LARS": cfg.SOLVER.LARS_ON, - }, - ] - optim_params = [x for x in optim_params if len(x["params"])] - - # Check all parameters will be passed into optimizer. - assert len(list(model.parameters())) == len(non_bn_parameters) + len( - bn_parameters - ) + len(zero_parameters) + len( - no_grad_parameters - ), "parameter size does not match: {} + {} + {} + {} != {}".format( - len(non_bn_parameters), - len(bn_parameters), - len(zero_parameters), - len(no_grad_parameters), - len(list(model.parameters())), - ) - print( - "bn {}, non bn {}, zero {} no grad {}".format( - len(bn_parameters), - len(non_bn_parameters), - len(zero_parameters), - len(no_grad_parameters), + raise ValueError( + "Layer decay should be in (0, 1], but is {}".format(cfg.SOLVER.LAYER_DECAY) ) - ) if cfg.SOLVER.OPTIMIZING_METHOD == "sgd": optimizer = torch.optim.SGD( @@ -105,13 +115,14 @@ def construct_optimizer(model, cfg): optimizer = torch.optim.Adam( optim_params, lr=cfg.SOLVER.BASE_LR, - betas=(0.9, 0.999), + betas=cfg.SOLVER.BETAS, weight_decay=cfg.SOLVER.WEIGHT_DECAY, ) elif cfg.SOLVER.OPTIMIZING_METHOD == "adamw": optimizer = torch.optim.AdamW( optim_params, lr=cfg.SOLVER.BASE_LR, + betas=cfg.SOLVER.BETAS, eps=1e-08, weight_decay=cfg.SOLVER.WEIGHT_DECAY, ) @@ -119,6 +130,7 @@ def construct_optimizer(model, cfg): optimizer = torch.optim._multi_tensor.AdamW( optim_params, lr=cfg.SOLVER.BASE_LR, + betas=cfg.SOLVER.BETAS, eps=1e-08, weight_decay=cfg.SOLVER.WEIGHT_DECAY, ) @@ -127,12 +139,104 @@ def construct_optimizer(model, cfg): "Does not support {} optimizer".format(cfg.SOLVER.OPTIMIZING_METHOD) ) if cfg.SOLVER.LARS_ON: - optimizer = LARS( - optimizer=optimizer, trust_coefficient=0.001, clip=False - ) + optimizer = LARS(optimizer=optimizer, trust_coefficient=0.001, clip=False) return optimizer +def get_param_groups(model, cfg): + def _get_layer_decay(name): + layer_id = None + if name in ("cls_token", "mask_token"): + layer_id = 0 + elif name.startswith("pos_embed"): + layer_id = 0 + elif name.startswith("patch_embed"): + layer_id = 0 + elif name.startswith("blocks"): + layer_id = int(name.split(".")[1]) + 1 + else: + layer_id = cfg.MVIT.DEPTH + 1 + layer_decay = cfg.SOLVER.LAYER_DECAY ** (cfg.MVIT.DEPTH + 1 - layer_id) + return layer_id, layer_decay + + for m in model.modules(): + assert not isinstance(m, torch.nn.modules.batchnorm._NormBase), ( + "BN is not supported with layer decay" + ) + + non_bn_parameters_count = 0 + zero_parameters_count = 0 + no_grad_parameters_count = 0 + parameter_group_names = {} + parameter_group_vars = {} + + skip = {} + if cfg.NUM_GPUS > 1: + if hasattr(model.module, "no_weight_decay"): + skip = model.module.no_weight_decay() + # skip = {"module." + v for v in skip} + else: + if hasattr(model, "no_weight_decay"): + skip = model.no_weight_decay() + + for name, p in model.named_parameters(): + if not p.requires_grad: + group_name = "no_grad" + no_grad_parameters_count += 1 + continue + name = name[len("module.") :] if name.startswith("module.") else name + if name in skip or ( + (len(p.shape) == 1 or name.endswith(".bias")) + and cfg.SOLVER.ZERO_WD_1D_PARAM + ): + layer_id, layer_decay = _get_layer_decay(name) + group_name = "layer_%d_%s" % (layer_id, "zero") + weight_decay = 0.0 + zero_parameters_count += 1 + else: + layer_id, layer_decay = _get_layer_decay(name) + group_name = "layer_%d_%s" % (layer_id, "non_bn") + weight_decay = cfg.SOLVER.WEIGHT_DECAY + non_bn_parameters_count += 1 + + if group_name not in parameter_group_names: + parameter_group_names[group_name] = { + "weight_decay": weight_decay, + "params": [], + "layer_decay": layer_decay, + } + parameter_group_vars[group_name] = { + "weight_decay": weight_decay, + "params": [], + "layer_decay": layer_decay, + } + parameter_group_names[group_name]["params"].append(name) + parameter_group_vars[group_name]["params"].append(p) + + # print("Param groups = %s" % json.dumps(parameter_group_names, indent=2)) + optim_params = list(parameter_group_vars.values()) + + # Check all parameters will be passed into optimizer. + assert ( + len(list(model.parameters())) + == non_bn_parameters_count + zero_parameters_count + no_grad_parameters_count + ), "parameter size does not match: {} + {} + {} != {}".format( + non_bn_parameters_count, + zero_parameters_count, + no_grad_parameters_count, + len(list(model.parameters())), + ) + print( + "non bn {}, zero {}, no grad {}".format( + non_bn_parameters_count, + zero_parameters_count, + no_grad_parameters_count, + ) + ) + + return optim_params + + def get_epoch_lr(cur_epoch, cfg): """ Retrieves the lr for the given epoch (as specified by the lr policy). @@ -152,10 +256,10 @@ def set_lr(optimizer, new_lr): new_lr (float): the new learning rate to set. """ for param_group in optimizer.param_groups: - param_group["lr"] = new_lr + param_group["lr"] = new_lr * param_group["layer_decay"] -class LARS(object): +class LARS: """ this class is adapted from https://github.com/NVIDIA/apex/blob/master/apex/parallel/LARC.py to include ignoring LARS application specific parameters (e.g. 1D params) @@ -219,13 +323,9 @@ def step(self): weight_decays = [] for group in self.optim.param_groups: # absorb weight decay control from optimizer - weight_decay = ( - group["weight_decay"] if "weight_decay" in group else 0 - ) + weight_decay = group["weight_decay"] if "weight_decay" in group else 0 weight_decays.append(weight_decay) - apply_LARS = ( - group["apply_LARS"] if "apply_LARS" in group else True - ) + apply_LARS = group["apply_LARS"] if "apply_LARS" in group else True if not apply_LARS: continue group["weight_decay"] = 0 @@ -257,3 +357,23 @@ def step(self): # return weight decay control to optimizer for i, group in enumerate(self.optim.param_groups): group["weight_decay"] = weight_decays[i] + + +def get_grad_norm_(parameters, norm_type=2.0): + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = [p for p in parameters if p.grad is not None] + norm_type = float(norm_type) + if len(parameters) == 0: + return torch.tensor(0.0) + device = parameters[0].grad.device + if norm_type == "inf": + total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters) + else: + total_norm = torch.norm( + torch.stack( + [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters] + ), + norm_type, + ) + return total_norm diff --git a/slowfast/models/ptv_model_builder.py b/slowfast/models/ptv_model_builder.py index 192dca59e..dd16c7fde 100644 --- a/slowfast/models/ptv_model_builder.py +++ b/slowfast/models/ptv_model_builder.py @@ -5,17 +5,11 @@ """Video models using PyTorchVideo model builder.""" from functools import partial + import torch.nn as nn from detectron2.layers import ROIAlign - -from slowfast.models.batchnorm_helper import get_norm -from slowfast.models.video_model_builder import _POOL1, _TEMPORAL_KERNEL_BASIS - from pytorchvideo.models.csn import create_csn -from pytorchvideo.models.head import ( - create_res_basic_head, - create_res_roi_pooling_head, -) +from pytorchvideo.models.head import create_res_basic_head, create_res_roi_pooling_head from pytorchvideo.models.r2plus1d import ( create_2plus1d_bottleneck_block, create_r2plus1d, @@ -25,11 +19,9 @@ from pytorchvideo.models.vision_transformers import ( create_multiscale_vision_transformers, ) -from pytorchvideo.models.x3d import ( - Swish, - create_x3d, - create_x3d_bottleneck_block, -) +from pytorchvideo.models.x3d import create_x3d, create_x3d_bottleneck_block, Swish +from slowfast.models.batchnorm_helper import get_norm +from slowfast.models.video_model_builder import _POOL1, _TEMPORAL_KERNEL_BASIS from .build import MODEL_REGISTRY @@ -50,8 +42,7 @@ def get_head_act(act_func): return nn.Sigmoid() else: raise NotImplementedError( - "{} is not supported as a head activation " - "function.".format(act_func) + "{} is not supported as a head activation function.".format(act_func) ) @@ -72,12 +63,10 @@ def __init__(self, cfg): """ super(PTVResNet, self).__init__() - assert ( - cfg.RESNET.STRIDE_1X1 is False - ), "STRIDE_1x1 must be True for PTVResNet" - assert ( - cfg.RESNET.TRANS_FUNC == "bottleneck_transform" - ), f"Unsupported TRANS_FUNC type {cfg.RESNET.TRANS_FUNC} for PTVResNet" + assert cfg.RESNET.STRIDE_1X1 is False, "STRIDE_1x1 must be True for PTVResNet" + assert cfg.RESNET.TRANS_FUNC == "bottleneck_transform", ( + f"Unsupported TRANS_FUNC type {cfg.RESNET.TRANS_FUNC} for PTVResNet" + ) assert cfg.MODEL.ARCH in [ "c2d", "slow", @@ -234,12 +223,10 @@ def __init__(self, cfg): """ super(PTVSlowFast, self).__init__() - assert ( - cfg.RESNET.STRIDE_1X1 is False - ), "STRIDE_1x1 must be True for PTVSlowFast" - assert ( - cfg.RESNET.TRANS_FUNC == "bottleneck_transform" - ), f"Unsupported TRANS_FUNC type {cfg.RESNET.TRANS_FUNC} for PTVSlowFast" + assert cfg.RESNET.STRIDE_1X1 is False, "STRIDE_1x1 must be True for PTVSlowFast" + assert cfg.RESNET.TRANS_FUNC == "bottleneck_transform", ( + f"Unsupported TRANS_FUNC type {cfg.RESNET.TRANS_FUNC} for PTVSlowFast" + ) self.detection_mode = cfg.DETECTION.ENABLE self._construct_network(cfg) @@ -272,10 +259,7 @@ def _construct_network(self, cfg): ((temp_kernel[stage + 1][pathway][0], 1, 1),) * num_block_temp_kernel[stage][pathway] + ((1, 1, 1),) - * ( - stage_depth[stage] - - num_block_temp_kernel[stage][pathway] - ) + * (stage_depth[stage] - num_block_temp_kernel[stage][pathway]) ) # Head from config @@ -300,9 +284,7 @@ def _construct_network(self, cfg): ) head_pool_kernel_sizes = ( ( - cfg.DATA.NUM_FRAMES - // cfg.SLOWFAST.ALPHA - // pool_size[0][0], + cfg.DATA.NUM_FRAMES // cfg.SLOWFAST.ALPHA // pool_size[0][0], 1, 1, ), @@ -311,9 +293,7 @@ def _construct_network(self, cfg): else: head_pool_kernel_sizes = ( ( - cfg.DATA.NUM_FRAMES - // cfg.SLOWFAST.ALPHA - // pool_size[0][0], + cfg.DATA.NUM_FRAMES // cfg.SLOWFAST.ALPHA // pool_size[0][0], cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][1], cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][2], ), @@ -434,15 +414,13 @@ def __init__(self, cfg): """ super(PTVX3D, self).__init__() - assert ( - cfg.RESNET.STRIDE_1X1 is False - ), "STRIDE_1x1 must be True for PTVX3D" - assert ( - cfg.RESNET.TRANS_FUNC == "x3d_transform" - ), f"Unsupported TRANS_FUNC type {cfg.RESNET.TRANS_FUNC} for PTVX3D" - assert ( - cfg.DETECTION.ENABLE is False - ), "Detection model is not supported for PTVX3D yet." + assert cfg.RESNET.STRIDE_1X1 is False, "STRIDE_1x1 must be True for PTVX3D" + assert cfg.RESNET.TRANS_FUNC == "x3d_transform", ( + f"Unsupported TRANS_FUNC type {cfg.RESNET.TRANS_FUNC} for PTVX3D" + ) + assert cfg.DETECTION.ENABLE is False, ( + "Detection model is not supported for PTVX3D yet." + ) self._construct_network(cfg) @@ -531,9 +509,9 @@ def __init__(self, cfg): """ super(PTVCSN, self).__init__() - assert ( - cfg.DETECTION.ENABLE is False - ), "Detection model is not supported for PTVCSN yet." + assert cfg.DETECTION.ENABLE is False, ( + "Detection model is not supported for PTVCSN yet." + ) self._construct_network(cfg) @@ -616,9 +594,9 @@ def __init__(self, cfg): """ super(PTVR2plus1D, self).__init__() - assert ( - cfg.DETECTION.ENABLE is False - ), "Detection model is not supported for PTVR2plus1D yet." + assert cfg.DETECTION.ENABLE is False, ( + "Detection model is not supported for PTVR2plus1D yet." + ) self._construct_network(cfg) @@ -717,9 +695,9 @@ def __init__(self, cfg): """ super(PTVMViT, self).__init__() - assert ( - cfg.DETECTION.ENABLE is False - ), "Detection model is not supported for PTVMViT yet." + assert cfg.DETECTION.ENABLE is False, ( + "Detection model is not supported for PTVMViT yet." + ) self._construct_network(cfg) diff --git a/slowfast/models/resnet_helper.py b/slowfast/models/resnet_helper.py index 63bc39bdd..449b114ec 100644 --- a/slowfast/models/resnet_helper.py +++ b/slowfast/models/resnet_helper.py @@ -3,9 +3,7 @@ """Video models.""" -import torch import torch.nn as nn - from slowfast.models.common import drop_path from slowfast.models.nonlocal_helper import Nonlocal from slowfast.models.operators import SE, Swish @@ -20,9 +18,9 @@ def get_trans_func(name): "basic_transform": BasicTransform, "x3d_transform": X3DTransform, } - assert ( - name in trans_funcs.keys() - ), "Transformation function '{}' not supported".format(name) + assert name in trans_funcs.keys(), ( + "Transformation function '{}' not supported".format(name) + ) return trans_funcs[name] @@ -694,9 +692,7 @@ def _construct( instantiation=instantiation, norm_module=norm_module, ) - self.add_module( - "pathway{}_nonlocal{}".format(pathway, i), nln - ) + self.add_module("pathway{}_nonlocal{}".format(pathway, i), nln) def forward(self, inputs): output = [] @@ -706,9 +702,7 @@ def forward(self, inputs): m = getattr(self, "pathway{}_res{}".format(pathway, i)) x = m(x) if hasattr(self, "pathway{}_nonlocal{}".format(pathway, i)): - nln = getattr( - self, "pathway{}_nonlocal{}".format(pathway, i) - ) + nln = getattr(self, "pathway{}_nonlocal{}".format(pathway, i)) b, c, t, h, w = x.shape if self.nonlocal_group[pathway] > 1: # Fold temporal dimension into batch dimension. diff --git a/slowfast/models/reversible_mvit.py b/slowfast/models/reversible_mvit.py new file mode 100644 index 000000000..bd5229a8a --- /dev/null +++ b/slowfast/models/reversible_mvit.py @@ -0,0 +1,672 @@ +import sys +from functools import partial + +import torch +from slowfast.models.attention import attention_pool, MultiScaleAttention +from slowfast.models.common import drop_path, Mlp, TwoStreamFusion +from slowfast.models.utils import round_width +from torch import nn +from torch.autograd import Function as Function + + +class ReversibleMViT(nn.Module): + """ + Reversible model builder. This builds the reversible transformer encoder + and allows reversible training. + + Karttikeya Mangalam, Haoqi Fan, Yanghao Li, Chao-Yuan Wu, Bo Xiong, + Christoph Feichtenhofer, Jitendra Malik + "Reversible Vision Transformers" + + https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf + """ + + def __init__(self, config, model): + """ + The `__init__` method of any subclass should also contain these + arguments. + Args: + cfg (CfgNode): model building configs, details are in the + comments of the config file. + model (nn.Module): parent MViT module this module forms + a reversible encoder in. + """ + + super().__init__() + self.cfg = config + + embed_dim = self.cfg.MVIT.EMBED_DIM + depth = self.cfg.MVIT.DEPTH + num_heads = self.cfg.MVIT.NUM_HEADS + mlp_ratio = self.cfg.MVIT.MLP_RATIO + qkv_bias = self.cfg.MVIT.QKV_BIAS + + drop_path_rate = self.cfg.MVIT.DROPPATH_RATE + self.dropout = config.MVIT.DROPOUT_RATE + self.pre_q_fusion = self.cfg.MVIT.REV.PRE_Q_FUSION + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + + input_size = model.patch_dims + + self.layers = nn.ModuleList([]) + self.no_custom_backward = False + + if self.cfg.MVIT.NORM == "layernorm": + norm_layer = partial(nn.LayerNorm, eps=1e-6) + else: + raise NotImplementedError("Only supports layernorm.") + + dim_mul, head_mul = torch.ones(depth + 1), torch.ones(depth + 1) + for i in range(len(self.cfg.MVIT.DIM_MUL)): + dim_mul[self.cfg.MVIT.DIM_MUL[i][0]] = self.cfg.MVIT.DIM_MUL[i][1] + for i in range(len(self.cfg.MVIT.HEAD_MUL)): + head_mul[self.cfg.MVIT.HEAD_MUL[i][0]] = self.cfg.MVIT.HEAD_MUL[i][1] + + pool_q = model.pool_q + pool_kv = model.pool_kv + stride_q = model.stride_q + stride_kv = model.stride_kv + + for i in range(depth): + num_heads = round_width(num_heads, head_mul[i]) + + # Upsampling inside the MHPA, input to the Q-pooling block is lower C dimension + # This localizes the feature changes in a single block, making more computation reversible. + embed_dim = round_width( + embed_dim, dim_mul[i - 1] if i > 0 else 1.0, divisor=num_heads + ) + dim_out = round_width( + embed_dim, + dim_mul[i], + divisor=round_width(num_heads, head_mul[i + 1]), + ) + + if i in self.cfg.MVIT.REV.BUFFER_LAYERS: + layer_type = StageTransitionBlock + input_mult = 2 if "concat" in self.pre_q_fusion else 1 + else: + layer_type = ReversibleBlock + input_mult = 1 + + dimout_correction = ( + 2 if (input_mult == 2 and "concat" in self.pre_q_fusion) else 1 + ) + + self.layers.append( + layer_type( + dim=embed_dim + * input_mult, # added only for concat fusion before Qpooling layers + input_size=input_size, + dim_out=dim_out * input_mult // dimout_correction, + num_heads=num_heads, + cfg=self.cfg, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + kernel_q=pool_q[i] if len(pool_q) > i else [], + kernel_kv=pool_kv[i] if len(pool_kv) > i else [], + stride_q=stride_q[i] if len(stride_q) > i else [], + stride_kv=stride_kv[i] if len(stride_kv) > i else [], + layer_id=i, + pre_q_fusion=self.pre_q_fusion, + ) + ) + # F is the attention block + self.layers[-1].F.thw = input_size + + if len(stride_q[i]) > 0: + input_size = [ + size // stride for size, stride in zip(input_size, stride_q[i]) + ] + + embed_dim = dim_out + + @staticmethod + def vanilla_backward(h, layers, buffer): + """ + Using rev layers without rev backpropagation. Debugging purposes only. + Activated with self.no_custom_backward. + """ + + # split into hidden states (h) and attention_output (a) + h, a = torch.chunk(h, 2, dim=-1) + for _, layer in enumerate(layers): + a, h = layer(a, h) + + return torch.cat([a, h], dim=-1) + + def forward(self, x): + # process the layers in a reversible stack and an irreversible stack. + stack = [] + for l_i in range(len(self.layers)): + if isinstance(self.layers[l_i], StageTransitionBlock): + stack.append(("StageTransition", l_i)) + else: + if len(stack) == 0 or stack[-1][0] == "StageTransition": + stack.append(("Reversible", [])) + stack[-1][1].append(l_i) + + for layer_seq in stack: + if layer_seq[0] == "StageTransition": + x = self.layers[layer_seq[1]](x) + + else: + x = torch.cat([x, x], dim=-1) + + # no need for custom backprop in eval/model stat log + if not self.training or self.no_custom_backward: + executing_fn = ReversibleMViT.vanilla_backward + else: + executing_fn = RevBackProp.apply + + x = executing_fn( + x, + self.layers[layer_seq[1][0] : layer_seq[1][-1] + 1], + [], # buffer activations + ) + + # Apply dropout + x = nn.functional.dropout(x, p=self.dropout, training=self.training) + + return x + + +class RevBackProp(Function): + """ + Custom Backpropagation function to allow (A) flusing memory in foward + and (B) activation recomputation reversibly in backward for gradient calculation. + + Inspired by https://github.com/RobinBruegger/RevTorch/blob/master/revtorch/revtorch.py + """ + + @staticmethod + def forward( + ctx, + x, + layers, + buffer_layers, # List of layer ids for int activation to buffer + ): + """ + Reversible Forward pass. Any intermediate activations from `buffer_layers` are + cached in ctx for forward pass. This is not necessary for standard usecases. + Each reversible layer implements its own forward pass logic. + """ + buffer_layers.sort() + + X_1, X_2 = torch.chunk(x, 2, dim=-1) + + intermediate = [] + + for layer in layers: + X_1, X_2 = layer(X_1, X_2) + + if layer.layer_id in buffer_layers: + intermediate.extend([X_1.detach(), X_2.detach()]) + + if len(buffer_layers) == 0: + all_tensors = [X_1.detach(), X_2.detach()] + else: + intermediate = [torch.LongTensor(buffer_layers), *intermediate] + all_tensors = [X_1.detach(), X_2.detach(), *intermediate] + + ctx.save_for_backward(*all_tensors) + ctx.layers = layers + + return torch.cat([X_1, X_2], dim=-1) + + @staticmethod + def backward(ctx, dx): + """ + Reversible Backward pass. Any intermediate activations from `buffer_layers` are + recovered from ctx. Each layer implements its own loic for backward pass (both + activation recomputation and grad calculation). + """ + dX_1, dX_2 = torch.chunk(dx, 2, dim=-1) + + # retrieve params from ctx for backward + X_1, X_2, *int_tensors = ctx.saved_tensors + + # no buffering + if len(int_tensors) != 0: + buffer_layers = int_tensors[0].tolist() + + else: + buffer_layers = [] + + layers = ctx.layers + + for _, layer in enumerate(layers[::-1]): + if layer.layer_id in buffer_layers: + X_1, X_2, dX_1, dX_2 = layer.backward_pass( + Y_1=int_tensors[buffer_layers.index(layer.layer_id) * 2 + 1], + Y_2=int_tensors[buffer_layers.index(layer.layer_id) * 2 + 2], + dY_1=dX_1, + dY_2=dX_2, + ) + + else: + X_1, X_2, dX_1, dX_2 = layer.backward_pass( + Y_1=X_1, + Y_2=X_2, + dY_1=dX_1, + dY_2=dX_2, + ) + + dx = torch.cat([dX_1, dX_2], dim=-1) + + del int_tensors + del dX_1, dX_2, X_1, X_2 + + return dx, None, None + + +class StageTransitionBlock(nn.Module): + """ + Blocks for changing the feature dimensions in MViT (using Q-pooling). + See Section 3.3.1 in paper for details. + """ + + def __init__( + self, + dim, + input_size, + dim_out, + num_heads, + mlp_ratio, + qkv_bias, + drop_path, + kernel_q, + kernel_kv, + stride_q, + stride_kv, + cfg, + norm_layer=nn.LayerNorm, + pre_q_fusion=None, + layer_id=0, + ): + """ + Uses the same structure of F and G functions as Reversible Block except + without using reversible forward (and backward) pass. + """ + super().__init__() + + self.drop_path_rate = drop_path + + embed_dim = dim + + self.F = AttentionSubBlock( + dim=embed_dim, + input_size=input_size, + num_heads=num_heads, + cfg=cfg, + dim_out=dim_out, + kernel_q=kernel_q, + kernel_kv=kernel_kv, + stride_q=stride_q, + stride_kv=stride_kv, + norm_layer=norm_layer, + ) + + self.G = MLPSubblock( + dim=dim_out, + mlp_ratio=mlp_ratio, + norm_layer=norm_layer, + ) + + self.layer_id = layer_id + + self.is_proj = False + self.has_cls_embed = cfg.MVIT.CLS_EMBED_ON + + self.is_conv = False + self.pool_first = cfg.MVIT.POOL_FIRST + self.mode = cfg.MVIT.MODE + self.pre_q_fuse = TwoStreamFusion(pre_q_fusion, dim=dim) + + if cfg.MVIT.REV.RES_PATH == "max": + self.res_conv = False + self.pool_skip = nn.MaxPool3d( + # self.attention.attn.pool_q.kernel_size, + [s + 1 if s > 1 else s for s in self.F.attn.pool_q.stride], + self.F.attn.pool_q.stride, + [int(k // 2) for k in self.F.attn.pool_q.stride], + # self.attention.attn.pool_q.padding, + ceil_mode=False, + ) + + elif cfg.MVIT.REV.RES_PATH == "conv": + self.res_conv = True + else: + raise NotImplementedError + + # Add a linear projection in residual branch + if embed_dim != dim_out: + self.is_proj = True + self.res_proj = nn.Linear(embed_dim, dim_out, bias=True) + + def forward( + self, + x, + ): + """ + Forward logic is similar to MultiScaleBlock with Q-pooling. + """ + x = self.pre_q_fuse(x) + + # fork tensor for residual connections + x_res = x + + # This uses conv to pool the residual hidden features + # but done before pooling only if not pool_first + if self.is_proj and not self.pool_first: + x_res = self.res_proj(x_res) + + if self.res_conv: + # Pooling the hidden features with the same conv as Q + N, L, C = x_res.shape + + # This handling is the same as that of q in MultiScaleAttention + if self.mode == "conv_unshared": + fold_dim = 1 + else: + fold_dim = self.F.attn.num_heads + + # Output is (B, N, L, C) + x_res = x_res.reshape(N, L, fold_dim, C // fold_dim).permute(0, 2, 1, 3) + + x_res, _ = attention_pool( + x_res, + self.F.attn.pool_q, + # thw_shape = self.attention.attn.thw, + thw_shape=self.F.thw, + has_cls_embed=self.has_cls_embed, + norm=self.F.attn.norm_q if hasattr(self.F.attn, "norm_q") else None, + ) + x_res = x_res.permute(0, 2, 1, 3).reshape(N, x_res.shape[2], C) + + else: + # Pooling the hidden features with max op + x_res, _ = attention_pool( + x_res, + self.pool_skip, + thw_shape=self.F.attn.thw, + has_cls_embed=self.has_cls_embed, + ) + + # If pool_first then project to higher dim now + if self.is_proj and self.pool_first: + x_res = self.res_proj(x_res) + + x = self.F(x) + x = x_res + x + x = x + self.G(x) + + x = drop_path(x, drop_prob=self.drop_path_rate, training=self.training) + + return x + + +class ReversibleBlock(nn.Module): + """ + Reversible Blocks for Reversible Vision Transformer and also + for state-preserving blocks in Reversible MViT. See Section + 3.3.2 in paper for details. + """ + + def __init__( + self, + dim, + input_size, + dim_out, + num_heads, + mlp_ratio, + qkv_bias, + drop_path, + kernel_q, + kernel_kv, + stride_q, + stride_kv, + cfg, + norm_layer=nn.LayerNorm, + layer_id=0, + **kwargs, + ): + """ + Block is composed entirely of function F (Attention + sub-block) and G (MLP sub-block) including layernorm. + """ + super().__init__() + + self.drop_path_rate = drop_path + + self.F = AttentionSubBlock( + dim=dim, + input_size=input_size, + num_heads=num_heads, + cfg=cfg, + dim_out=dim_out, + kernel_q=kernel_q, + kernel_kv=kernel_kv, + stride_q=stride_q, + stride_kv=stride_kv, + norm_layer=norm_layer, + ) + + self.G = MLPSubblock( + dim=dim, + mlp_ratio=mlp_ratio, + norm_layer=norm_layer, + ) + + self.layer_id = layer_id + + self.seeds = {} + + def seed_cuda(self, key): + """ + Fix seeds to allow for stochastic elements such as + dropout to be reproduced exactly in activation + recomputation in the backward pass. + """ + + # randomize seeds + # use cuda generator if available + if ( + hasattr(torch.cuda, "default_generators") + and len(torch.cuda.default_generators) > 0 + ): + # GPU + device_idx = torch.cuda.current_device() + seed = torch.cuda.default_generators[device_idx].seed() + else: + # CPU + seed = int(torch.seed() % sys.maxsize) + + self.seeds[key] = seed + torch.manual_seed(self.seeds[key]) + + def forward(self, X_1, X_2): + """ + forward pass equations: + Y_1 = X_1 + Attention(X_2), F = Attention + Y_2 = X_2 + MLP(Y_1), G = MLP + """ + + self.seed_cuda("attn") + # Y_1 : attn_output + f_X_2 = self.F(X_2) + + self.seed_cuda("droppath") + f_X_2_dropped = drop_path( + f_X_2, drop_prob=self.drop_path_rate, training=self.training + ) + + # Y_1 = X_1 + f(X_2) + Y_1 = X_1 + f_X_2_dropped + + # free memory + del X_1 + + self.seed_cuda("FFN") + g_Y_1 = self.G(Y_1) + + torch.manual_seed(self.seeds["droppath"]) + g_Y_1_dropped = drop_path( + g_Y_1, drop_prob=self.drop_path_rate, training=self.training + ) + + # Y_2 = X_2 + g(Y_1) + Y_2 = X_2 + g_Y_1_dropped + + del X_2 + + return Y_1, Y_2 + + def backward_pass( + self, + Y_1, + Y_2, + dY_1, + dY_2, + ): + """ + equation for activation recomputation: + X_2 = Y_2 - G(Y_1), G = MLP + X_1 = Y_1 - F(X_2), F = Attention + """ + + # temporarily record intermediate activation for G + # and use them for gradient calculcation of G + with torch.enable_grad(): + Y_1.requires_grad = True + + torch.manual_seed(self.seeds["FFN"]) + g_Y_1 = self.G(Y_1) + + torch.manual_seed(self.seeds["droppath"]) + g_Y_1 = drop_path( + g_Y_1, drop_prob=self.drop_path_rate, training=self.training + ) + + g_Y_1.backward(dY_2, retain_graph=True) + + # activation recomputation is by design and not part of + # the computation graph in forward pass. + with torch.no_grad(): + X_2 = Y_2 - g_Y_1 + del g_Y_1 + + dY_1 = dY_1 + Y_1.grad + Y_1.grad = None + + # record F activations and calc gradients on F + with torch.enable_grad(): + X_2.requires_grad = True + + torch.manual_seed(self.seeds["attn"]) + f_X_2 = self.F(X_2) + + torch.manual_seed(self.seeds["droppath"]) + f_X_2 = drop_path( + f_X_2, drop_prob=self.drop_path_rate, training=self.training + ) + + f_X_2.backward(dY_1, retain_graph=True) + + # propagate reverse computed acitvations at the start of + # the previou block for backprop.s + with torch.no_grad(): + X_1 = Y_1 - f_X_2 + + del f_X_2, Y_1 + dY_2 = dY_2 + X_2.grad + + X_2.grad = None + X_2 = X_2.detach() + + return X_1, X_2, dY_1, dY_2 + + +class MLPSubblock(nn.Module): + """ + This creates the function G such that the entire block can be + expressed as F(G(X)). Includes pre-LayerNorm. + """ + + def __init__( + self, + dim, + mlp_ratio, + norm_layer=nn.LayerNorm, + ): + super().__init__() + self.norm = norm_layer(dim, eps=1e-6, elementwise_affine=True) + + mlp_hidden_dim = int(dim * mlp_ratio) + + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=nn.GELU, + ) + + def forward(self, x): + return self.mlp(self.norm(x)) + + +class AttentionSubBlock(nn.Module): + """ + This creates the function F such that the entire block can be + expressed as F(G(X)). Includes pre-LayerNorm. + """ + + def __init__( + self, + dim, + input_size, + num_heads, + cfg, + dim_out=None, + kernel_q=(1, 1, 1), + kernel_kv=(1, 1, 1), + stride_q=(1, 1, 1), + stride_kv=(1, 1, 1), + norm_layer=nn.LayerNorm, + ): + super().__init__() + self.norm = norm_layer(dim, eps=1e-6, elementwise_affine=True) + + # This will be set externally during init + self.thw = None + + # the actual attention details are the same as Multiscale + # attention for MViTv2 (with channel up=projection inside block) + # can also implement no upprojection attention for vanilla ViT + self.attn = MultiScaleAttention( + dim, + dim_out, + input_size=input_size, + num_heads=num_heads, + kernel_q=kernel_q, + kernel_kv=kernel_kv, + stride_q=stride_q, + stride_kv=stride_kv, + norm_layer=norm_layer, + drop_rate=cfg.MVIT.DROPOUT_RATE, + qkv_bias=cfg.MVIT.QKV_BIAS, + has_cls_embed=cfg.MVIT.CLS_EMBED_ON, + mode=cfg.MVIT.MODE, + pool_first=cfg.MVIT.POOL_FIRST, + rel_pos_spatial=cfg.MVIT.REL_POS_SPATIAL, + rel_pos_temporal=cfg.MVIT.REL_POS_TEMPORAL, + rel_pos_zero_init=cfg.MVIT.REL_POS_ZERO_INIT, + residual_pooling=cfg.MVIT.RESIDUAL_POOLING, + separate_qkv=cfg.MVIT.SEPARATE_QKV, + ) + + def forward(self, x): + out, _ = self.attn(self.norm(x), self.thw) + return out diff --git a/slowfast/models/stem_helper.py b/slowfast/models/stem_helper.py index 889961317..a34607779 100644 --- a/slowfast/models/stem_helper.py +++ b/slowfast/models/stem_helper.py @@ -11,9 +11,9 @@ def get_stem_func(name): Retrieves the stem module by name. """ trans_funcs = {"x3d_stem": X3DStem, "basic_stem": ResNetBasicStem} - assert ( - name in trans_funcs.keys() - ), "Transformation function '{}' not supported".format(name) + assert name in trans_funcs.keys(), ( + "Transformation function '{}' not supported".format(name) + ) return trans_funcs[name] @@ -113,9 +113,9 @@ def _construct_stem(self, dim_in, dim_out, norm_module, stem_func_name): self.add_module("pathway{}_stem".format(pathway), stem) def forward(self, x): - assert ( - len(x) == self.num_pathways - ), "Input tensor does not contain {} pathway".format(self.num_pathways) + assert len(x) == self.num_pathways, ( + "Input tensor does not contain {} pathway".format(self.num_pathways) + ) # use a new list, don't modify in-place the x list, which is bad for activation checkpointing. y = [] for pathway in range(len(x)): @@ -187,9 +187,7 @@ def _construct_stem(self, dim_in, dim_out, norm_module): padding=self.padding, bias=False, ) - self.bn = norm_module( - num_features=dim_out, eps=self.eps, momentum=self.bn_mmt - ) + self.bn = norm_module(num_features=dim_out, eps=self.eps, momentum=self.bn_mmt) self.relu = nn.ReLU(self.inplace_relu) self.pool_layer = nn.MaxPool3d( kernel_size=[1, 3, 3], stride=[1, 2, 2], padding=[0, 1, 1] @@ -276,9 +274,7 @@ def _construct_stem(self, dim_in, dim_out, norm_module): groups=dim_out, ) - self.bn = norm_module( - num_features=dim_out, eps=self.eps, momentum=self.bn_mmt - ) + self.bn = norm_module(num_features=dim_out, eps=self.eps, momentum=self.bn_mmt) self.relu = nn.ReLU(self.inplace_relu) def forward(self, x): @@ -316,7 +312,9 @@ def __init__( padding=padding, ) - def forward(self, x): + def forward(self, x, keep_spatial=False): x = self.proj(x) + if keep_spatial: + return x, x.shape # B C (T) H W -> B (T)HW C return x.flatten(2).transpose(1, 2), x.shape diff --git a/slowfast/models/utils.py b/slowfast/models/utils.py index 381ba98fc..a3014a5b8 100644 --- a/slowfast/models/utils.py +++ b/slowfast/models/utils.py @@ -1,6 +1,8 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +import numpy as np import slowfast.utils.logging as logging +import torch logger = logging.get_logger(__name__) @@ -27,3 +29,184 @@ def validate_checkpoint_wrapper_import(checkpoint_wrapper): """ if checkpoint_wrapper is None: raise ImportError("Please install fairscale.") + + +def get_gkern(kernlen, std): + """Returns a 2D Gaussian kernel array.""" + + def _gaussian_fn(kernlen, std): + n = torch.arange(0, kernlen).float() + n -= n.mean() + n /= std + w = torch.exp(-0.5 * n**2) + return w + + gkern1d = _gaussian_fn(kernlen, std) + gkern2d = torch.outer(gkern1d, gkern1d) + return gkern2d / gkern2d.sum() + + +# -------------------------------------------------------- +# 2D sine-cosine position embedding +# References: +# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py +# MoCo v3: https://github.com/facebookresearch/moco-v3 +# -------------------------------------------------------- +def get_3d_sincos_pos_embed(embed_dim, grid_size, t_size, cls_token=False): + """ + grid_size: int of the grid height and width + t_size: int of the temporal size + return: + pos_embed: [t_size*grid_size*grid_size, embed_dim] or [1+t_size*grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + assert embed_dim % 4 == 0 + embed_dim_spatial = embed_dim // 4 * 3 + embed_dim_temporal = embed_dim // 4 + + # spatial + grid_h = np.arange(grid_size, dtype=np.float32) + grid_w = np.arange(grid_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size, grid_size]) + pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid) + + # temporal + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t) + + # concate: [T, H, W] order + pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :] + pos_embed_temporal = np.repeat( + pos_embed_temporal, grid_size**2, axis=1 + ) # [T, H*W, D // 4] + pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :] + pos_embed_spatial = np.repeat( + pos_embed_spatial, t_size, axis=0 + ) # [T, H*W, D // 4 * 3] + + pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1) + pos_embed = pos_embed.reshape([-1, embed_dim]) # [T*H*W, D] + + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): + """ + grid_size: int of the grid height and width + return: + pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + grid_h = np.arange(grid_size, dtype=np.float32) + grid_w = np.arange(grid_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size, grid_size]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + assert embed_dim % 2 == 0 + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) + + emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) + return emb + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +# -------------------------------------------------------- +# Interpolate position embeddings for high-resolution +# References: +# DeiT: https://github.com/facebookresearch/deit +# -------------------------------------------------------- +def interpolate_pos_embed(model, checkpoint_model): + if "pos_embed" in checkpoint_model: + pos_embed_checkpoint = checkpoint_model["pos_embed"] + embedding_size = pos_embed_checkpoint.shape[-1] + num_patches = model.patch_embed.num_patches + num_extra_tokens = model.pos_embed.shape[-2] - num_patches + # height (== width) for the checkpoint position embedding + orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5) + # height (== width) for the new position embedding + new_size = int(num_patches**0.5) + # class_token and dist_token are kept unchanged + if orig_size != new_size: + print( + "Position interpolate from %dx%d to %dx%d" + % (orig_size, orig_size, new_size, new_size) + ) + extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] + # only the position tokens are interpolated + pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] + pos_tokens = pos_tokens.reshape( + -1, orig_size, orig_size, embedding_size + ).permute(0, 3, 1, 2) + pos_tokens = torch.nn.functional.interpolate( + pos_tokens, + size=(new_size, new_size), + mode="bicubic", + align_corners=False, + ) + pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) + new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) + checkpoint_model["pos_embed"] = new_pos_embed + + +def calc_mvit_feature_geometry(cfg): + feat_size = [ + [ + ( + cfg.DATA.NUM_FRAMES // cfg.MVIT.PATCH_STRIDE[0] + if len(cfg.MVIT.PATCH_STRIDE) > 2 + else 1 + ), + cfg.DATA.TRAIN_CROP_SIZE // cfg.MVIT.PATCH_STRIDE[-2], + cfg.DATA.TRAIN_CROP_SIZE // cfg.MVIT.PATCH_STRIDE[-1], + ] + for i in range(cfg.MVIT.DEPTH) + ] + feat_stride = [ + [ + cfg.MVIT.PATCH_STRIDE[0] if len(cfg.MVIT.PATCH_STRIDE) > 2 else 1, + cfg.MVIT.PATCH_STRIDE[-2], + cfg.MVIT.PATCH_STRIDE[-1], + ] + for i in range(cfg.MVIT.DEPTH) + ] + for _, x in enumerate(cfg.MVIT.POOL_Q_STRIDE): + for i in range(cfg.MVIT.DEPTH): + if i >= x[0]: + for j in range(len(feat_size[i])): + feat_size[i][j] = feat_size[i][j] // x[j + 1] + feat_stride[i][j] = feat_stride[i][j] * x[j + 1] + return feat_size, feat_stride diff --git a/slowfast/models/video_model_builder.py b/slowfast/models/video_model_builder.py index 44b20db5e..e0f39ab8f 100644 --- a/slowfast/models/video_model_builder.py +++ b/slowfast/models/video_model_builder.py @@ -5,20 +5,25 @@ import math from functools import partial + +import slowfast.utils.logging as logging +import slowfast.utils.weight_init_helper as init_helper import torch import torch.nn as nn import torch.nn.functional as F -from torch.nn.init import trunc_normal_ - -import slowfast.utils.weight_init_helper as init_helper from slowfast.models.attention import MultiScaleBlock from slowfast.models.batchnorm_helper import get_norm +from slowfast.models.common import TwoStreamFusion +from slowfast.models.reversible_mvit import ReversibleMViT from slowfast.models.utils import ( + calc_mvit_feature_geometry, + get_3d_sincos_pos_embed, round_width, validate_checkpoint_wrapper_import, ) +from torch.nn.init import trunc_normal_ -from . import head_helper, resnet_helper, stem_helper +from . import head_helper, operators, resnet_helper, stem_helper # noqa from .build import MODEL_REGISTRY try: @@ -27,6 +32,8 @@ checkpoint_wrapper = None +logger = logging.get_logger(__name__) + # Number of blocks for different stages given the model depth. _MODEL_STAGE_DEPTH = {18: (2, 2, 2, 2), 50: (3, 4, 6, 3), 101: (3, 4, 23, 3)} @@ -211,9 +218,7 @@ def _construct_network(self, cfg): num_groups = cfg.RESNET.NUM_GROUPS width_per_group = cfg.RESNET.WIDTH_PER_GROUP dim_inner = num_groups * width_per_group - out_dim_ratio = ( - cfg.SLOWFAST.BETA_INV // cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO - ) + out_dim_ratio = cfg.SLOWFAST.BETA_INV // cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH] @@ -370,9 +375,7 @@ def _construct_network(self, cfg): num_classes=cfg.MODEL.NUM_CLASSES, pool_size=[ [ - cfg.DATA.NUM_FRAMES - // cfg.SLOWFAST.ALPHA - // pool_size[0][0], + cfg.DATA.NUM_FRAMES // cfg.SLOWFAST.ALPHA // pool_size[0][0], 1, 1, ], @@ -392,23 +395,25 @@ def _construct_network(self, cfg): width_per_group * 32 // cfg.SLOWFAST.BETA_INV, ], num_classes=cfg.MODEL.NUM_CLASSES, - pool_size=[None, None] - if cfg.MULTIGRID.SHORT_CYCLE - or cfg.MODEL.MODEL_NAME == "ContrastiveModel" - else [ - [ - cfg.DATA.NUM_FRAMES - // cfg.SLOWFAST.ALPHA - // pool_size[0][0], - cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][1], - cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][2], - ], - [ - cfg.DATA.NUM_FRAMES // pool_size[1][0], - cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[1][1], - cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[1][2], - ], - ], # None for AdaptiveAvgPool3d((1, 1, 1)) + pool_size=( + [None, None] + if cfg.MULTIGRID.SHORT_CYCLE + or cfg.MODEL.MODEL_NAME == "ContrastiveModel" + else [ + [ + cfg.DATA.NUM_FRAMES + // cfg.SLOWFAST.ALPHA + // pool_size[0][0], + cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][1], + cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][2], + ], + [ + cfg.DATA.NUM_FRAMES // pool_size[1][0], + cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[1][1], + cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[1][2], + ], + ] + ), # None for AdaptiveAvgPool3d((1, 1, 1)) dropout_rate=cfg.MODEL.DROPOUT_RATE, act_func=cfg.MODEL.HEAD_ACT, detach_final_fc=cfg.MODEL.DETACH_FINAL_FC, @@ -619,16 +624,18 @@ def _construct_network(self, cfg): self.head = head_helper.ResNetBasicHead( dim_in=[width_per_group * 32], num_classes=cfg.MODEL.NUM_CLASSES, - pool_size=[None] - if cfg.MULTIGRID.SHORT_CYCLE - or cfg.MODEL.MODEL_NAME == "ContrastiveModel" - else [ - [ - cfg.DATA.NUM_FRAMES // pool_size[0][0], - cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][1], - cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][2], + pool_size=( + [None] + if cfg.MULTIGRID.SHORT_CYCLE + or cfg.MODEL.MODEL_NAME == "ContrastiveModel" + else [ + [ + cfg.DATA.NUM_FRAMES // pool_size[0][0], + cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][1], + cfg.DATA.TRAIN_CROP_SIZE // 32 // pool_size[0][2], + ] ] - ], # None for AdaptiveAvgPool3d((1, 1, 1)) + ), # None for AdaptiveAvgPool3d((1, 1, 1)) dropout_rate=cfg.MODEL.DROPOUT_RATE, act_func=cfg.MODEL.HEAD_ACT, detach_final_fc=cfg.MODEL.DETACH_FINAL_FC, @@ -748,9 +755,7 @@ def _construct_network(self, cfg): dim_inner = int(cfg.X3D.BOTTLENECK_FACTOR * dim_out) n_rep = self._round_repeats(block[0], d_mul) - prefix = "s{}".format( - stage + 2 - ) # start w res2 to follow convention + prefix = "s{}".format(stage + 2) # start w res2 to follow convention s = resnet_helper.ResStage( dim_in=[dim_in], @@ -759,9 +764,7 @@ def _construct_network(self, cfg): temp_kernel_sizes=temp_kernel[1], stride=[block[2]], num_blocks=[n_rep], - num_groups=[dim_inner] - if cfg.X3D.CHANNELWISE_3x3x3 - else [num_groups], + num_groups=[dim_inner] if cfg.X3D.CHANNELWISE_3x3x3 else [num_groups], num_block_temp_kernel=[n_rep], nonlocal_inds=cfg.NONLOCAL.LOCATION[0], nonlocal_group=cfg.NONLOCAL.GROUP[0], @@ -822,10 +825,15 @@ def __init__(self, cfg): spatial_size = cfg.DATA.TRAIN_CROP_SIZE temporal_size = cfg.DATA.NUM_FRAMES in_chans = cfg.DATA.INPUT_CHANNEL_NUM[0] - use_2d_patch = cfg.MVIT.PATCH_2D + self.use_2d_patch = cfg.MVIT.PATCH_2D + self.enable_detection = cfg.DETECTION.ENABLE + self.enable_rev = cfg.MVIT.REV.ENABLE self.patch_stride = cfg.MVIT.PATCH_STRIDE - if use_2d_patch: + if self.use_2d_patch: self.patch_stride = [1] + self.patch_stride + self.T = cfg.DATA.NUM_FRAMES // self.patch_stride[0] + self.H = cfg.DATA.TRAIN_CROP_SIZE // self.patch_stride[1] + self.W = cfg.DATA.TRAIN_CROP_SIZE // self.patch_stride[2] # Prepare output. num_classes = cfg.MODEL.NUM_CLASSES embed_dim = cfg.MVIT.EMBED_DIM @@ -836,10 +844,14 @@ def __init__(self, cfg): self.drop_rate = cfg.MVIT.DROPOUT_RATE depth = cfg.MVIT.DEPTH drop_path_rate = cfg.MVIT.DROPPATH_RATE + layer_scale_init_value = cfg.MVIT.LAYER_SCALE_INIT_VALUE + head_init_scale = cfg.MVIT.HEAD_INIT_SCALE mode = cfg.MVIT.MODE self.cls_embed_on = cfg.MVIT.CLS_EMBED_ON + self.use_mean_pooling = cfg.MVIT.USE_MEAN_POOLING # Params for positional embedding self.use_abs_pos = cfg.MVIT.USE_ABS_POS + self.use_fixed_sincos_pos = cfg.MVIT.USE_FIXED_SINCOS_POS self.sep_pos_embed = cfg.MVIT.SEP_POS_EMBED self.rel_pos_spatial = cfg.MVIT.REL_POS_SPATIAL self.rel_pos_temporal = cfg.MVIT.REL_POS_TEMPORAL @@ -848,17 +860,17 @@ def __init__(self, cfg): else: raise NotImplementedError("Only supports layernorm.") self.num_classes = num_classes - patch_embed = stem_helper.PatchEmbed( + self.patch_embed = stem_helper.PatchEmbed( dim_in=in_chans, dim_out=embed_dim, kernel=cfg.MVIT.PATCH_KERNEL, stride=cfg.MVIT.PATCH_STRIDE, padding=cfg.MVIT.PATCH_PADDING, - conv_2d=use_2d_patch, + conv_2d=self.use_2d_patch, ) + if cfg.MODEL.ACT_CHECKPOINT: - patch_embed = checkpoint_wrapper(patch_embed) - self.patch_embed = patch_embed + self.patch_embed = checkpoint_wrapper(self.patch_embed) self.input_dims = [temporal_size, spatial_size, spatial_size] assert self.input_dims[1] == self.input_dims[2] self.patch_dims = [ @@ -880,20 +892,21 @@ def __init__(self, cfg): if self.use_abs_pos: if self.sep_pos_embed: self.pos_embed_spatial = nn.Parameter( - torch.zeros( - 1, self.patch_dims[1] * self.patch_dims[2], embed_dim - ) + torch.zeros(1, self.patch_dims[1] * self.patch_dims[2], embed_dim) ) self.pos_embed_temporal = nn.Parameter( torch.zeros(1, self.patch_dims[0], embed_dim) ) if self.cls_embed_on: - self.pos_embed_class = nn.Parameter( - torch.zeros(1, 1, embed_dim) - ) + self.pos_embed_class = nn.Parameter(torch.zeros(1, 1, embed_dim)) else: self.pos_embed = nn.Parameter( - torch.zeros(1, pos_embed_dim, embed_dim) + torch.zeros( + 1, + pos_embed_dim, + embed_dim, + ), + requires_grad=not self.use_fixed_sincos_pos, ) if self.drop_rate > 0.0: @@ -911,9 +924,7 @@ def __init__(self, cfg): stride_kv = [[] for i in range(cfg.MVIT.DEPTH)] for i in range(len(cfg.MVIT.POOL_Q_STRIDE)): - stride_q[cfg.MVIT.POOL_Q_STRIDE[i][0]] = cfg.MVIT.POOL_Q_STRIDE[i][ - 1: - ] + stride_q[cfg.MVIT.POOL_Q_STRIDE[i][0]] = cfg.MVIT.POOL_Q_STRIDE[i][1:] if cfg.MVIT.POOL_KVQ_KERNEL is not None: pool_q[cfg.MVIT.POOL_Q_STRIDE[i][0]] = cfg.MVIT.POOL_KVQ_KERNEL else: @@ -934,84 +945,115 @@ def __init__(self, cfg): cfg.MVIT.POOL_KV_STRIDE.append([i] + _stride_kv) for i in range(len(cfg.MVIT.POOL_KV_STRIDE)): - stride_kv[cfg.MVIT.POOL_KV_STRIDE[i][0]] = cfg.MVIT.POOL_KV_STRIDE[ - i - ][1:] + stride_kv[cfg.MVIT.POOL_KV_STRIDE[i][0]] = cfg.MVIT.POOL_KV_STRIDE[i][1:] if cfg.MVIT.POOL_KVQ_KERNEL is not None: - pool_kv[ - cfg.MVIT.POOL_KV_STRIDE[i][0] - ] = cfg.MVIT.POOL_KVQ_KERNEL + pool_kv[cfg.MVIT.POOL_KV_STRIDE[i][0]] = cfg.MVIT.POOL_KVQ_KERNEL else: pool_kv[cfg.MVIT.POOL_KV_STRIDE[i][0]] = [ - s + 1 if s > 1 else s - for s in cfg.MVIT.POOL_KV_STRIDE[i][1:] + s + 1 if s > 1 else s for s in cfg.MVIT.POOL_KV_STRIDE[i][1:] ] + self.pool_q = pool_q + self.pool_kv = pool_kv + self.stride_q = stride_q + self.stride_kv = stride_kv + self.norm_stem = norm_layer(embed_dim) if cfg.MVIT.NORM_STEM else None input_size = self.patch_dims - self.blocks = nn.ModuleList() - if cfg.MODEL.ACT_CHECKPOINT: - validate_checkpoint_wrapper_import(checkpoint_wrapper) + if self.enable_rev: + # rev does not allow cls token + assert not self.cls_embed_on - for i in range(depth): - num_heads = round_width(num_heads, head_mul[i]) - if cfg.MVIT.DIM_MUL_IN_ATT: - dim_out = round_width( - embed_dim, - dim_mul[i], - divisor=round_width(num_heads, head_mul[i]), - ) + self.rev_backbone = ReversibleMViT(cfg, self) + + embed_dim = round_width(embed_dim, dim_mul.prod(), divisor=num_heads) + + self.fuse = TwoStreamFusion(cfg.MVIT.REV.RESPATH_FUSE, dim=2 * embed_dim) + + if "concat" in self.cfg.MVIT.REV.RESPATH_FUSE: + self.norm = norm_layer(2 * embed_dim) else: - dim_out = round_width( - embed_dim, - dim_mul[i + 1], - divisor=round_width(num_heads, head_mul[i + 1]), + self.norm = norm_layer(embed_dim) + + else: + self.blocks = nn.ModuleList() + + for i in range(depth): + num_heads = round_width(num_heads, head_mul[i]) + if cfg.MVIT.DIM_MUL_IN_ATT: + dim_out = round_width( + embed_dim, + dim_mul[i], + divisor=round_width(num_heads, head_mul[i]), + ) + else: + dim_out = round_width( + embed_dim, + dim_mul[i + 1], + divisor=round_width(num_heads, head_mul[i + 1]), + ) + attention_block = MultiScaleBlock( + dim=embed_dim, + dim_out=dim_out, + num_heads=num_heads, + input_size=input_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop_rate=self.drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + kernel_q=pool_q[i] if len(pool_q) > i else [], + kernel_kv=pool_kv[i] if len(pool_kv) > i else [], + stride_q=stride_q[i] if len(stride_q) > i else [], + stride_kv=stride_kv[i] if len(stride_kv) > i else [], + mode=mode, + has_cls_embed=self.cls_embed_on, + pool_first=pool_first, + rel_pos_spatial=self.rel_pos_spatial, + rel_pos_temporal=self.rel_pos_temporal, + rel_pos_zero_init=cfg.MVIT.REL_POS_ZERO_INIT, + residual_pooling=cfg.MVIT.RESIDUAL_POOLING, + dim_mul_in_att=cfg.MVIT.DIM_MUL_IN_ATT, + separate_qkv=cfg.MVIT.SEPARATE_QKV, ) - attention_block = MultiScaleBlock( - dim=embed_dim, - dim_out=dim_out, - num_heads=num_heads, - input_size=input_size, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - drop_rate=self.drop_rate, - drop_path=dpr[i], - norm_layer=norm_layer, - kernel_q=pool_q[i] if len(pool_q) > i else [], - kernel_kv=pool_kv[i] if len(pool_kv) > i else [], - stride_q=stride_q[i] if len(stride_q) > i else [], - stride_kv=stride_kv[i] if len(stride_kv) > i else [], - mode=mode, - has_cls_embed=self.cls_embed_on, - pool_first=pool_first, - rel_pos_spatial=self.rel_pos_spatial, - rel_pos_temporal=self.rel_pos_temporal, - rel_pos_zero_init=cfg.MVIT.REL_POS_ZERO_INIT, - residual_pooling=cfg.MVIT.RESIDUAL_POOLING, - dim_mul_in_att=cfg.MVIT.DIM_MUL_IN_ATT, - separate_qkv=cfg.MVIT.SEPARATE_QKV, + + if cfg.MODEL.ACT_CHECKPOINT: + attention_block = checkpoint_wrapper(attention_block) + self.blocks.append(attention_block) + if len(stride_q[i]) > 0: + input_size = [ + size // stride for size, stride in zip(input_size, stride_q[i]) + ] + + embed_dim = dim_out + + self.norm = norm_layer(embed_dim) + + if self.enable_detection: + self.head = head_helper.ResNetRoIHead( + dim_in=[embed_dim], + num_classes=num_classes, + pool_size=[[temporal_size // self.patch_stride[0], 1, 1]], + resolution=[[cfg.DETECTION.ROI_XFORM_RESOLUTION] * 2], + scale_factor=[cfg.DETECTION.SPATIAL_SCALE_FACTOR], + dropout_rate=cfg.MODEL.DROPOUT_RATE, + act_func=cfg.MODEL.HEAD_ACT, + aligned=cfg.DETECTION.ALIGNED, + ) + else: + self.head = head_helper.TransformerBasicHead( + ( + 2 * embed_dim + if ("concat" in cfg.MVIT.REV.RESPATH_FUSE and self.enable_rev) + else embed_dim + ), + num_classes, + dropout_rate=cfg.MODEL.DROPOUT_RATE, + act_func=cfg.MODEL.HEAD_ACT, + cfg=cfg, ) - if cfg.MODEL.ACT_CHECKPOINT: - attention_block = checkpoint_wrapper(attention_block) - self.blocks.append(attention_block) - - if len(stride_q[i]) > 0: - input_size = [ - size // stride - for size, stride in zip(input_size, stride_q[i]) - ] - embed_dim = dim_out - self.norm = norm_layer(embed_dim) - - self.head = head_helper.TransformerBasicHead( - embed_dim, - num_classes, - dropout_rate=cfg.MODEL.DROPOUT_RATE, - act_func=cfg.MODEL.HEAD_ACT, - cfg=cfg, - ) if self.use_abs_pos: if self.sep_pos_embed: trunc_normal_(self.pos_embed_spatial, std=0.02) @@ -1020,17 +1062,33 @@ def __init__(self, cfg): trunc_normal_(self.pos_embed_class, std=0.02) else: trunc_normal_(self.pos_embed, std=0.02) + if self.use_fixed_sincos_pos: + pos_embed = get_3d_sincos_pos_embed( + self.pos_embed.shape[-1], + self.H, + self.T, + cls_token=self.cls_embed_on, + ) + self.pos_embed.data.copy_( + torch.from_numpy(pos_embed).float().unsqueeze(0) + ) + if self.cls_embed_on: trunc_normal_(self.cls_token, std=0.02) self.apply(self._init_weights) + self.head.projection.weight.data.mul_(head_init_scale) + self.head.projection.bias.data.mul_(head_init_scale) + + self.feat_size, self.feat_stride = calc_mvit_feature_geometry(cfg) + def _init_weights(self, m): - if isinstance(m, nn.Linear): + if isinstance(m, (nn.Linear, nn.Conv2d, nn.Conv3d)): nn.init.trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) + nn.init.constant_(m.bias, 0.02) elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) + nn.init.constant_(m.bias, 0.02) nn.init.constant_(m.weight, 1.0) @torch.jit.ignore @@ -1047,7 +1105,7 @@ def no_weight_decay(self): ] ) else: - names.append(["pos_embed"]) + names.append("pos_embed") if self.rel_pos_spatial: names.extend(["rel_pos_h", "rel_pos_w", "rel_pos_hw"]) if self.rel_pos_temporal: @@ -1058,7 +1116,10 @@ def no_weight_decay(self): return names def _get_pos_embed(self, pos_embed, bcthw): - t, h, w = bcthw[-3], bcthw[-2], bcthw[-1] + if len(bcthw) == 4: + t, h, w = 1, bcthw[-2], bcthw[-1] + else: + t, h, w = bcthw[-3], bcthw[-2], bcthw[-1] if self.cls_embed_on: cls_pos_embed = pos_embed[:, 0:1, :] pos_embed = pos_embed[:, 1:] @@ -1068,9 +1129,7 @@ def _get_pos_embed(self, pos_embed, bcthw): if (p_t, p_h, p_w) != (t, h, w): new_pos_embed = F.interpolate( - pos_embed[:, :, :] - .reshape(1, p_t, p_h, p_w, -1) - .permute(0, 4, 1, 2, 3), + pos_embed[:, :, :].reshape(1, p_t, p_h, p_w, -1).permute(0, 4, 1, 2, 3), size=(t, h, w), mode="trilinear", ) @@ -1081,18 +1140,49 @@ def _get_pos_embed(self, pos_embed, bcthw): return pos_embed - def forward(self, x): + def _forward_reversible(self, x): + """ + Reversible specific code for forward computation. + """ + # rev does not support cls token or detection + assert not self.cls_embed_on + assert not self.enable_detection + + x = self.rev_backbone(x) + + if self.use_mean_pooling: + x = self.fuse(x) + x = x.mean(1) + x = self.norm(x) + else: + x = self.norm(x) + x = self.fuse(x) + x = x.mean(1) + + x = self.head(x) + + return x + + def forward(self, x, bboxes=None, return_attn=False): x = x[0] x, bcthw = self.patch_embed(x) - - T = self.cfg.DATA.NUM_FRAMES // self.patch_stride[0] - H, W = bcthw[-2], bcthw[-1] + bcthw = list(bcthw) + if len(bcthw) == 4: # Fix bcthw in case of 4D tensor + bcthw.insert(2, torch.tensor(self.T)) + T, H, W = bcthw[-3], bcthw[-2], bcthw[-1] + assert len(bcthw) == 5 and (T, H, W) == (self.T, self.H, self.W), bcthw B, N, C = x.shape + s = 1 if self.cls_embed_on else 0 + if self.use_fixed_sincos_pos: + x += self.pos_embed[:, s:, :] # s: on/off cls token + if self.cls_embed_on: cls_tokens = self.cls_token.expand( B, -1, -1 ) # stole cls_tokens impl from Phil Wang, thanks + if self.use_fixed_sincos_pos: + cls_tokens = cls_tokens + self.pos_embed[:, :s, :] x = torch.cat((cls_tokens, x), dim=1) if self.use_abs_pos: @@ -1106,11 +1196,9 @@ def forward(self, x): ) if self.cls_embed_on: pos_embed = torch.cat([self.pos_embed_class, pos_embed], 1) - pos_embed = self._get_pos_embed(pos_embed, bcthw) - x = x + pos_embed + x += self._get_pos_embed(pos_embed, bcthw) else: - pos_embed = self._get_pos_embed(self.pos_embed, bcthw) - x = x + pos_embed + x += self._get_pos_embed(self.pos_embed, bcthw) if self.drop_rate: x = self.pos_drop(x) @@ -1119,14 +1207,38 @@ def forward(self, x): x = self.norm_stem(x) thw = [T, H, W] - for blk in self.blocks: - x, thw = blk(x, thw) - x = self.norm(x) - if self.cls_embed_on: - x = x[:, 0] + if self.enable_rev: + x = self._forward_reversible(x) + else: - x = x.mean(1) + for blk in self.blocks: + x, thw = blk(x, thw) + + if self.enable_detection: + assert not self.enable_rev + + x = self.norm(x) + if self.cls_embed_on: + x = x[:, 1:] + + B, _, C = x.shape + x = x.transpose(1, 2).reshape(B, C, thw[0], thw[1], thw[2]) + + x = self.head([x], bboxes) + + else: + if self.use_mean_pooling: + if self.cls_embed_on: + x = x[:, 1:] + x = x.mean(1) + x = self.norm(x) + elif self.cls_embed_on: + x = self.norm(x) + x = x[:, 0] + else: # this is default, [norm->mean] + x = self.norm(x) + x = x.mean(1) + x = self.head(x) - x = self.head(x) return x diff --git a/slowfast/utils/ava_eval_helper.py b/slowfast/utils/ava_eval_helper.py index 4fb92093a..e3a127760 100644 --- a/slowfast/utils/ava_eval_helper.py +++ b/slowfast/utils/ava_eval_helper.py @@ -23,25 +23,21 @@ """Helper functions for AVA evaluation.""" -from __future__ import ( - absolute_import, - division, - print_function, - unicode_literals, -) +from __future__ import absolute_import, division, print_function, unicode_literals + import csv import logging -import numpy as np import pprint import time from collections import defaultdict +import numpy as np import slowfast.utils.distributed as du -from slowfast.utils.ava_evaluation import ( +from slowfast.utils.env import pathmgr +from vision.fair.slowfast.ava_evaluation import ( object_detection_evaluation, standard_fields, ) -from slowfast.utils.env import pathmgr logger = logging.getLogger(__name__) @@ -158,9 +154,7 @@ def evaluate_ava( ) logger.info("Evaluating with %d unique GT frames." % len(groundtruth[0])) - logger.info( - "Evaluating with %d unique detection frames" % len(detections[0]) - ) + logger.info("Evaluating with %d unique detection frames" % len(detections[0])) write_results(detections, "detections_%s.csv" % name) write_results(groundtruth, "groundtruth_%s.csv" % name) @@ -171,14 +165,10 @@ def evaluate_ava( return results["PascalBoxes_Precision/mAP@0.5IOU"] -def run_evaluation( - categories, groundtruth, detections, excluded_keys, verbose=True -): +def run_evaluation(categories, groundtruth, detections, excluded_keys, verbose=True): """AVA evaluation main logic.""" - pascal_evaluator = object_detection_evaluation.PascalDetectionEvaluator( - categories - ) + pascal_evaluator = object_detection_evaluation.PascalDetectionEvaluator(categories) boxes, labels, _ = groundtruth @@ -188,10 +178,7 @@ def run_evaluation( for image_key in boxes: if image_key in excluded_keys: logging.info( - ( - "Found excluded timestamp in ground truth: %s. " - "It will be ignored." - ), + ("Found excluded timestamp in ground truth: %s. It will be ignored."), image_key, ) continue @@ -217,10 +204,7 @@ def run_evaluation( for image_key in boxes: if image_key in excluded_keys: logging.info( - ( - "Found excluded timestamp in detections: %s. " - "It will be ignored." - ), + ("Found excluded timestamp in detections: %s. It will be ignored."), image_key, ) continue diff --git a/slowfast/utils/ava_evaluation/__init__.py b/slowfast/utils/ava_evaluation/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/slowfast/utils/benchmark.py b/slowfast/utils/benchmark.py old mode 100755 new mode 100644 index 33e5fe907..da04fcea4 --- a/slowfast/utils/benchmark.py +++ b/slowfast/utils/benchmark.py @@ -3,14 +3,14 @@ Functions for benchmarks. """ -import numpy as np import pprint -import torch -import tqdm -from fvcore.common.timer import Timer +import numpy as np import slowfast.utils.logging as logging import slowfast.utils.misc as misc +import torch +import tqdm +from fvcore.common.timer import Timer from slowfast.datasets import loader from slowfast.utils.env import setup_environment @@ -40,9 +40,7 @@ def benchmark_data_loading(cfg): timer = Timer() dataloader = loader.construct_loader(cfg, "train") - logger.info( - "Initialize loader using {:.2f} seconds.".format(timer.seconds()) - ) + logger.info("Initialize loader using {:.2f} seconds.".format(timer.seconds())) # Total batch size across different machines. batch_size = cfg.TRAIN.BATCH_SIZE * cfg.NUM_SHARDS log_period = cfg.BENCHMARK.LOG_PERIOD diff --git a/slowfast/utils/bn_helper.py b/slowfast/utils/bn_helper.py index da81c81e1..7e8b11f06 100644 --- a/slowfast/utils/bn_helper.py +++ b/slowfast/utils/bn_helper.py @@ -4,6 +4,7 @@ """bn helper.""" import itertools + import torch @@ -49,9 +50,7 @@ def compute_and_update_bn_stats(model, data_loader, num_batches=200): running_mean = [torch.zeros_like(bn.running_mean) for bn in bn_layers] running_square_mean = [torch.zeros_like(bn.running_var) for bn in bn_layers] - for ind, (inputs, _, _) in enumerate( - itertools.islice(data_loader, num_batches) - ): + for ind, (inputs, _, _) in enumerate(itertools.islice(data_loader, num_batches)): # Forwards the model to update the bn stats. if isinstance(inputs, (list,)): for i in range(len(inputs)): @@ -65,9 +64,9 @@ def compute_and_update_bn_stats(model, data_loader, num_batches=200): running_mean[i] += (bn.running_mean - running_mean[i]) / (ind + 1) # $E(x^2) = Var(x) + E(x)^2$. cur_square_mean = bn.running_var + bn.running_mean**2 - running_square_mean[i] += ( - cur_square_mean - running_square_mean[i] - ) / (ind + 1) + running_square_mean[i] += (cur_square_mean - running_square_mean[i]) / ( + ind + 1 + ) for i, bn in enumerate(bn_layers): bn.running_mean = running_mean[i] diff --git a/slowfast/utils/checkpoint.py b/slowfast/utils/checkpoint.py index 33176d4b5..e47620231 100644 --- a/slowfast/utils/checkpoint.py +++ b/slowfast/utils/checkpoint.py @@ -4,14 +4,15 @@ """Functions that handle saving and loading of checkpoints.""" import copy -import numpy as np +import math import os import pickle from collections import OrderedDict -import torch +import numpy as np import slowfast.utils.distributed as du import slowfast.utils.logging as logging +import torch from slowfast.utils.c2_model_loading import get_name_convert_func from slowfast.utils.env import checkpoint_pathmgr as pathmgr @@ -102,9 +103,7 @@ def is_checkpoint_epoch(cfg, cur_epoch, multigrid_schedule=None): prev_epoch = 0 for s in multigrid_schedule: if cur_epoch < s[-1]: - period = max( - (s[-1] - prev_epoch) // cfg.MULTIGRID.EVAL_FREQ + 1, 1 - ) + period = max((s[-1] - prev_epoch) // cfg.MULTIGRID.EVAL_FREQ + 1, 1) return (s[-1] - 1 - cur_epoch) % period == 0 prev_epoch = s[-1] @@ -140,9 +139,7 @@ def save_checkpoint(path_to_job, model, optimizer, epoch, cfg, scaler=None): if scaler is not None: checkpoint["scaler_state"] = scaler.state_dict() # Write the checkpoint. - path_to_checkpoint = get_path_to_checkpoint( - path_to_job, epoch + 1, cfg.TASK - ) + path_to_checkpoint = get_path_to_checkpoint(path_to_job, epoch + 1, cfg.TASK) with pathmgr.open(path_to_checkpoint, "wb") as f: torch.save(checkpoint, f) return path_to_checkpoint @@ -166,22 +163,16 @@ def inflate_weight(state_dict_2d, state_dict_3d): v3d = state_dict_3d[k] # Inflate the weight of 2D conv to 3D conv. if len(v2d.shape) == 4 and len(v3d.shape) == 5: - logger.info( - "Inflate {}: {} -> {}: {}".format(k, v2d.shape, k, v3d.shape) - ) + logger.info("Inflate {}: {} -> {}: {}".format(k, v2d.shape, k, v3d.shape)) # Dimension need to be match. assert v2d.shape[-2:] == v3d.shape[-2:] assert v2d.shape[:2] == v3d.shape[:2] - v3d = ( - v2d.unsqueeze(2).repeat(1, 1, v3d.shape[2], 1, 1) / v3d.shape[2] - ) + v3d = v2d.unsqueeze(2).repeat(1, 1, v3d.shape[2], 1, 1) / v3d.shape[2] elif v2d.shape == v3d.shape: v3d = v2d else: logger.info( - "Unexpected {}: {} -|> {}: {}".format( - k, v2d.shape, k, v3d.shape - ) + "Unexpected {}: {} -|> {}: {}".format(k, v2d.shape, k, v3d.shape) ) state_dict_inflated[k] = v3d.clone() return state_dict_inflated @@ -197,6 +188,7 @@ def load_checkpoint( convert_from_caffe2=False, epoch_reset=False, clear_name_pattern=(), + image_init=False, ): """ Load the checkpoint from the given file. If inflation is True, inflate the @@ -235,9 +227,7 @@ def load_checkpoint( # expand shape dims if they differ (eg for converting linear to conv params) if len(c2_blob_shape) < len(model_blob_shape): - c2_blob_shape += (1,) * ( - len(model_blob_shape) - len(c2_blob_shape) - ) + c2_blob_shape += (1,) * (len(model_blob_shape) - len(c2_blob_shape)) caffe2_checkpoint["blobs"][key] = np.reshape( caffe2_checkpoint["blobs"][key], c2_blob_shape ) @@ -280,9 +270,7 @@ def load_checkpoint( prefix in key for prefix in ["momentum", "lr", "model_iter"] ): logger.warn( - "!! {}: can not be converted, got {}".format( - key, converted_key - ) + "!! {}: can not be converted, got {}".format(key, converted_key) ) diff = set(ms.state_dict()) - set(state_dict) diff = {d for d in diff if "num_batches_tracked" not in d} @@ -315,36 +303,212 @@ def load_checkpoint( k_re = k.replace( item, "", 1 ) # only repace first occurence of pattern - model_state_dict_new[k_re] = checkpoint[ - "model_state" - ][k] + model_state_dict_new[k_re] = checkpoint["model_state"][k] logger.info("renaming: {} -> {}".format(k, k_re)) else: - model_state_dict_new[k] = checkpoint["model_state"][ - k - ] + model_state_dict_new[k] = checkpoint["model_state"][k] checkpoint["model_state"] = model_state_dict_new pre_train_dict = checkpoint["model_state"] model_dict = ms.state_dict() + + if image_init: + if ( + "pos_embed" in pre_train_dict.keys() + and "pos_embed_xy" in model_dict.keys() + ): + print( + pre_train_dict["pos_embed"].shape, + model_dict["pos_embed_xy"].shape, + model_dict["pos_embed_class"].shape, + ) + if ( + pre_train_dict["pos_embed"].shape[1] + == model_dict["pos_embed_xy"].shape[1] + 1 + ): + pre_train_dict["pos_embed_xy"] = pre_train_dict["pos_embed"][ + :, 1: + ] + pre_train_dict["pos_embed_class"] = pre_train_dict["pos_embed"][ + :, :1 + ] + + if ( + "patch_embed.proj.weight" in pre_train_dict.keys() + and "patch_embed.proj.weight" in model_dict.keys() + ): + print( + pre_train_dict["patch_embed.proj.weight"].shape, + model_dict["patch_embed.proj.weight"].shape, + ) + if ( + len(pre_train_dict["patch_embed.proj.weight"].shape) == 4 + and len(model_dict["patch_embed.proj.weight"].shape) == 5 + ): # img->video + t = model_dict["patch_embed.proj.weight"].shape[2] + pre_train_dict["patch_embed.proj.weight"] = pre_train_dict[ + "patch_embed.proj.weight" + ][:, :, None, :, :].repeat(1, 1, t, 1, 1) + logger.info( + f"inflate patch_embed.proj.weight to {pre_train_dict['patch_embed.proj.weight'].shape}" + ) + elif ( + len(pre_train_dict["patch_embed.proj.weight"].shape) == 5 + and len(model_dict["patch_embed.proj.weight"].shape) == 4 + ): # video->img + orig_shape = pre_train_dict["patch_embed.proj.weight"].shape + # pre_train_dict["patch_embed.proj.weight"] = pre_train_dict["patch_embed.proj.weight"][:, :, orig_shape[2]//2, :, :] # take center + pre_train_dict["patch_embed.proj.weight"] = pre_train_dict[ + "patch_embed.proj.weight" + ].sum(2) # take avg + logger.info( + f"deflate patch_embed.proj.weight from {orig_shape} to {pre_train_dict['patch_embed.proj.weight'].shape}" + ) + if ( + "pos_embed_spatial" in pre_train_dict.keys() + and "pos_embed" in model_dict.keys() + ): + pos_embds = pre_train_dict["pos_embed_spatial"] + if ( + "pos_embed_class" in pre_train_dict.keys() + and pos_embds.shape != model_dict["pos_embed"].shape + ): + pos_embds = torch.cat( + [ + pre_train_dict["pos_embed_class"], + pos_embds, + ], + 1, + ) + pre_train_dict.pop("pos_embed_class") + if pos_embds.shape == model_dict["pos_embed"].shape: + pre_train_dict["pos_embed"] = pos_embds + pre_train_dict.pop("pos_embed_spatial") + logger.info( + f"successful surgery of pos embed w/ shape {pos_embds.shape} " + ) + else: + logger.info( + f"UNSUCCESSFUL surgery of pos embed w/ shape {pos_embds.shape} " + ) + + qkv = [ + "attn.pool_k.weight", + "attn.pool_q.weight", + "attn.pool_v.weight", + ] + for k in pre_train_dict.keys(): + if ( + any([x in k for x in qkv]) + and pre_train_dict[k].shape != model_dict[k].shape + ): + # print(pre_train_dict[k].shape, model_dict[k].shape) + logger.info( + f"inflate {k} from {pre_train_dict[k].shape} to {model_dict[k].shape}" + ) + t = model_dict[k].shape[2] + pre_train_dict[k] = pre_train_dict[k].repeat(1, 1, t, 1, 1) + + for k in pre_train_dict.keys(): + if ( + "rel_pos" in k + and pre_train_dict[k].shape != model_dict[k].shape + ): + # print(pre_train_dict[k].shape, model_dict[k].shape) + logger.info( + f"interpolating {k} from {pre_train_dict[k].shape} to {model_dict[k].shape}" + ) + new_pos_embed = torch.nn.functional.interpolate( + pre_train_dict[k] + .reshape(1, pre_train_dict[k].shape[0], -1) + .permute(0, 2, 1), + size=model_dict[k].shape[0], + mode="linear", + ) + new_pos_embed = ( + new_pos_embed.reshape(-1, model_dict[k].shape[0]) + .permute(1, 0) + .squeeze() + ) + pre_train_dict[k] = new_pos_embed + # Match pre-trained weights that have same shape as current model. - pre_train_dict_match = { - k: v - for k, v in pre_train_dict.items() - if k in model_dict and v.size() == model_dict[k].size() - } + pre_train_dict_match = {} + not_used_layers = [] + for k, v in pre_train_dict.items(): + if k in model_dict: + if v.size() == model_dict[k].size(): + pre_train_dict_match[k] = v + else: + if "attn.rel_pos" in k: + v_shape = v.shape + v = v.t().unsqueeze(0) + v = torch.nn.functional.interpolate( + v, + size=model_dict[k].size()[0], + mode="linear", + ) + v = v[0].t() + pre_train_dict_match[k] = v + logger.info( + "{} reshaped from {} to {}".format(k, v_shape, v.shape) + ) + elif "pos_embed_temporal" in k: + v_shape = v.shape + v = torch.nn.functional.interpolate( + v.permute(0, 2, 1), + size=model_dict[k].shape[1], + mode="linear", + ) + pre_train_dict_match[k] = v.permute(0, 2, 1) + logger.info( + "{} reshaped from {} to {}".format( + k, v_shape, pre_train_dict_match[k].shape + ) + ) + elif "pos_embed_spatial" in k: + v_shape = v.shape + pretrain_size = int(math.sqrt(v_shape[1])) + model_size = int(math.sqrt(model_dict[k].shape[1])) + assert pretrain_size * pretrain_size == v_shape[1] + assert model_size * model_size == model_dict[k].shape[1] + v = torch.nn.functional.interpolate( + v.reshape(1, pretrain_size, pretrain_size, -1).permute( + 0, 3, 1, 2 + ), + size=(model_size, model_size), + mode="bicubic", + ) + pre_train_dict_match[k] = v.reshape( + 1, -1, model_size * model_size + ).permute(0, 2, 1) + logger.info( + "{} reshaped from {} to {}".format( + k, v_shape, pre_train_dict_match[k].shape + ) + ) + else: + not_used_layers.append(k) + else: + not_used_layers.append(k) # Weights that do not have match from the pre-trained model. not_load_layers = [ - k - for k in model_dict.keys() - if k not in pre_train_dict_match.keys() + k for k in model_dict.keys() if k not in pre_train_dict_match.keys() ] # Log weights that are not loaded with the pre-trained weights. if not_load_layers: for k in not_load_layers: logger.info("Network weights {} not loaded.".format(k)) + if not_used_layers: + for k in not_used_layers: + logger.info("Network weights {} not used.".format(k)) # Load pre-trained weights. - ms.load_state_dict(pre_train_dict_match, strict=False) + missing_keys, unexpected_keys = ms.load_state_dict( + pre_train_dict_match, strict=False + ) + + print("missing keys: {}".format(missing_keys)) + print("unexpected keys: {}".format(unexpected_keys)) epoch = -1 # Load the optimizer state (commonly not done when fine-tuning) @@ -450,13 +614,10 @@ def normal_to_sub_bn(checkpoint_sd, model_sd): ): before_shape = checkpoint_sd[key].shape checkpoint_sd[key] = torch.cat( - [checkpoint_sd[key]] - * (model_blob_shape[0] // c2_blob_shape[0]) + [checkpoint_sd[key]] * (model_blob_shape[0] // c2_blob_shape[0]) ) logger.info( - "{} {} -> {}".format( - key, before_shape, checkpoint_sd[key].shape - ) + "{} {} -> {}".format(key, before_shape, checkpoint_sd[key].shape) ) return checkpoint_sd @@ -527,6 +688,7 @@ def load_train_checkpoint(cfg, model, optimizer, scaler=None): convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2", epoch_reset=cfg.TRAIN.CHECKPOINT_EPOCH_RESET, clear_name_pattern=cfg.TRAIN.CHECKPOINT_CLEAR_NAME_PATTERN, + image_init=cfg.TRAIN.CHECKPOINT_IN_INIT, ) start_epoch = checkpoint_epoch + 1 else: diff --git a/slowfast/utils/distributed.py b/slowfast/utils/distributed.py index b7de50842..5a6472d09 100644 --- a/slowfast/utils/distributed.py +++ b/slowfast/utils/distributed.py @@ -6,19 +6,23 @@ import functools import logging import pickle + import torch import torch.distributed as dist - from pytorchvideo.layers.distributed import ( # noqa cat_all_gather, get_local_process_group, get_local_rank, get_local_size, get_world_size, - init_distributed_training, + init_distributed_training as _init_distributed_training, ) +def init_distributed_training(cfg): + return _init_distributed_training(cfg.NUM_GPUS, cfg.SHARD_ID) + + def all_gather(tensors): """ All gathers the provided tensors from all processes across machines. @@ -31,9 +35,7 @@ def all_gather(tensors): output_tensor = [] world_size = dist.get_world_size() for tensor in tensors: - tensor_placeholder = [ - torch.ones_like(tensor) for _ in range(world_size) - ] + tensor_placeholder = [torch.ones_like(tensor) for _ in range(world_size)] dist.all_gather(tensor_placeholder, tensor, async_op=False) gather_list.append(tensor_placeholder) for gathered_tensor in gather_list: @@ -197,12 +199,10 @@ def _pad_to_largest_tensor(tensor, group): Tensor: padded tensor that has the max size """ world_size = dist.get_world_size(group=group) - assert ( - world_size >= 1 - ), "comm.gather/all_gather must be called from ranks within the given group!" - local_size = torch.tensor( - [tensor.numel()], dtype=torch.int64, device=tensor.device + assert world_size >= 1, ( + "comm.gather/all_gather must be called from ranks within the given group!" ) + local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device) size_list = [ torch.zeros([1], dtype=torch.int64, device=tensor.device) for _ in range(world_size) @@ -292,7 +292,6 @@ def forward(ctx, input): @staticmethod def backward(ctx, grad_output): - reduction = torch.distributed.all_reduce(grad_output, async_op=True) reduction.wait() diff --git a/slowfast/utils/logging.py b/slowfast/utils/logging.py index a19db4960..e899edb8f 100644 --- a/slowfast/utils/logging.py +++ b/slowfast/utils/logging.py @@ -10,8 +10,8 @@ import logging import os import sys -import simplejson +import simplejson import slowfast.utils.distributed as du from slowfast.utils.env import pathmgr @@ -30,9 +30,7 @@ def print_pass(*objects, sep=" ", end="\n", file=sys.stdout, flush=False): @functools.lru_cache(maxsize=None) def _cached_log_stream(filename): # Use 1K buffer if writing to cloud storage. - io = pathmgr.open( - filename, "a", buffering=1024 if "://" in filename else -1 - ) + io = pathmgr.open(filename, "a", buffering=1024 if "://" in filename else -1) atexit.register(io.close) return io @@ -91,13 +89,13 @@ def log_json_stats(stats, output_dir=None): stats (dict): a dictionary of statistical information to log. """ stats = { - k: decimal.Decimal("{:.3f}".format(v)) if isinstance(v, float) else v + k: decimal.Decimal("{:.5f}".format(v)) if isinstance(v, float) else v for k, v in stats.items() } json_stats = simplejson.dumps(stats, sort_keys=True, use_decimal=True) logger = get_logger(__name__) logger.info("json_stats: {:s}".format(json_stats)) - if du.is_root_proc() and output_dir: + if du.is_master_proc(du.get_world_size()) and output_dir: filename = os.path.join(output_dir, "json_stats.log") try: with pathmgr.open( @@ -105,6 +103,4 @@ def log_json_stats(stats, output_dir=None): ) as f: f.write("json_stats: {:s}\n".format(json_stats)) except Exception: - logger.info( - "Failed to write to json_stats.log: {}".format(json_stats) - ) + logger.info("Failed to write to json_stats.log: {}".format(json_stats)) diff --git a/slowfast/utils/lr_policy.py b/slowfast/utils/lr_policy.py index 80cfdd4c2..85b700656 100644 --- a/slowfast/utils/lr_policy.py +++ b/slowfast/utils/lr_policy.py @@ -19,9 +19,7 @@ def get_lr_at_epoch(cfg, cur_epoch): # Perform warm up. if cur_epoch < cfg.SOLVER.WARMUP_EPOCHS: lr_start = cfg.SOLVER.WARMUP_START_LR - lr_end = get_lr_func(cfg.SOLVER.LR_POLICY)( - cfg, cfg.SOLVER.WARMUP_EPOCHS - ) + lr_end = get_lr_func(cfg.SOLVER.LR_POLICY)(cfg, cfg.SOLVER.WARMUP_EPOCHS) alpha = (lr_end - lr_start) / cfg.SOLVER.WARMUP_EPOCHS lr = cur_epoch * alpha + lr_start return lr @@ -44,9 +42,7 @@ def lr_func_cosine(cfg, cur_epoch): cfg.SOLVER.COSINE_END_LR + (cfg.SOLVER.BASE_LR - cfg.SOLVER.COSINE_END_LR) * ( - math.cos( - math.pi * (cur_epoch - offset) / (cfg.SOLVER.MAX_EPOCH - offset) - ) + math.cos(math.pi * (cur_epoch - offset) / (cfg.SOLVER.MAX_EPOCH - offset)) + 1.0 ) * 0.5 diff --git a/slowfast/utils/meters.py b/slowfast/utils/meters.py index 48ff1ef47..cccf06d8c 100644 --- a/slowfast/utils/meters.py +++ b/slowfast/utils/meters.py @@ -4,17 +4,17 @@ """Meters.""" import datetime -import numpy as np import os from collections import defaultdict, deque -import torch -from fvcore.common.timer import Timer -from sklearn.metrics import average_precision_score +import numpy as np import slowfast.datasets.ava_helper as ava_helper import slowfast.utils.logging as logging import slowfast.utils.metrics as metrics import slowfast.utils.misc as misc +import torch +from fvcore.common.timer import Timer +from sklearn.metrics import average_precision_score from slowfast.utils.ava_eval_helper import ( evaluate_ava, read_csv, @@ -43,7 +43,7 @@ def get_ava_mini_groundtruth(full_groundtruth): return ret -class AVAMeter(object): +class AVAMeter: """ Measure the AVA train, val, and test stats. """ @@ -72,15 +72,11 @@ def __init__(self, overall_iters, cfg, mode): self.categories, self.class_whitelist = read_labelmap( os.path.join(cfg.AVA.ANNOTATION_DIR, cfg.AVA.LABEL_MAP_FILE) ) - gt_filename = os.path.join( - cfg.AVA.ANNOTATION_DIR, cfg.AVA.GROUNDTRUTH_FILE - ) + gt_filename = os.path.join(cfg.AVA.ANNOTATION_DIR, cfg.AVA.GROUNDTRUTH_FILE) self.full_groundtruth = read_csv(gt_filename, self.class_whitelist) self.mini_groundtruth = get_ava_mini_groundtruth(self.full_groundtruth) - _, self.video_idx_to_name = ava_helper.load_image_lists( - cfg, mode == "train" - ) + _, self.video_idx_to_name = ava_helper.load_image_lists(cfg, mode == "train") self.output_dir = cfg.OUTPUT_DIR self.min_top1_err = 100.0 @@ -105,9 +101,7 @@ def log_iter_stats(self, cur_epoch, cur_iter): if self.mode == "train": stats = { "_type": "{}_iter".format(self.mode), - "cur_epoch": "{}/{}".format( - cur_epoch + 1, self.cfg.SOLVER.MAX_EPOCH - ), + "cur_epoch": "{}/{}".format(cur_epoch + 1, self.cfg.SOLVER.MAX_EPOCH), "cur_iter": "{}".format(cur_iter + 1), "eta": eta, "dt": self.iter_timer.seconds(), @@ -120,9 +114,7 @@ def log_iter_stats(self, cur_epoch, cur_iter): elif self.mode == "val": stats = { "_type": "{}_iter".format(self.mode), - "cur_epoch": "{}/{}".format( - cur_epoch + 1, self.cfg.SOLVER.MAX_EPOCH - ), + "cur_epoch": "{}/{}".format(cur_epoch + 1, self.cfg.SOLVER.MAX_EPOCH), "cur_iter": "{}".format(cur_iter + 1), "eta": eta, "dt": self.iter_timer.seconds(), @@ -244,7 +236,7 @@ def log_epoch_stats(self, cur_epoch): logging.log_json_stats(stats, self.output_dir) -class TestMeter(object): +class TestMeter: """ Perform the multi-view ensemble for testing: each video with an unique index will be sampled with multiple clips, and the predictions of the clips will @@ -339,9 +331,7 @@ def update_stats(self, preds, labels, clip_ids): ) else: raise NotImplementedError( - "Ensemble Method {} is not supported".format( - self.ensemble_method - ) + "Ensemble Method {} is not supported".format(self.ensemble_method) ) self.clip_count[vid_id] += 1 @@ -408,20 +398,15 @@ def finalize_metrics(self, ks=(1, 5)): num_topks_correct = metrics.topks_correct( self.video_preds, self.video_labels, ks ) - topks = [ - (x / self.video_preds.size(0)) * 100.0 - for x in num_topks_correct - ] + topks = [(x / self.video_preds.size(0)) * 100.0 for x in num_topks_correct] assert len({len(ks), len(topks)}) == 1 for k, topk in zip(ks, topks): # self.stats["top{}_acc".format(k)] = topk.cpu().numpy() - self.stats["top{}_acc".format(k)] = "{:.{prec}f}".format( - topk, prec=2 - ) + self.stats["top{}_acc".format(k)] = "{:.{prec}f}".format(topk, prec=2) logging.log_json_stats(self.stats) -class ScalarMeter(object): +class ScalarMeter: """ A scalar meter uses a deque to track a series of scaler values with a given window size. It supports calculating the median and average values of the @@ -459,6 +444,9 @@ def get_win_median(self): """ return np.median(self.deque) + def get_current_value(self): + return self.deque[-1] + def get_win_avg(self): """ Calculate the current average value of the deque. @@ -472,7 +460,43 @@ def get_global_avg(self): return self.total / self.count -class TrainMeter(object): +class ListMeter: + def __init__(self, list_size): + """ + Args: + list_size (int): size of the list. + """ + self.list = np.zeros(list_size) + self.total = np.zeros(list_size) + self.count = 0 + + def reset(self): + """ + Reset the meter. + """ + self.list = np.zeros_like(self.list) + self.total = np.zeros_like(self.total) + self.count = 0 + + def add_value(self, value): + """ + Add a new list value to the meter. + """ + self.list = np.array(value) + self.count += 1 + self.total += self.list + + def get_value(self): + return self.list + + def get_global_avg(self): + """ + Calculate the global mean value. + """ + return self.total / self.count + + +class TrainMeter: """ Measure training stats. """ @@ -492,6 +516,7 @@ def __init__(self, epoch_iters, cfg): self.loss = ScalarMeter(cfg.LOG_PERIOD) self.loss_total = 0.0 self.lr = None + self.grad_norm = None # Current minibatch errors (smoothed over a window). self.mb_top1_err = ScalarMeter(cfg.LOG_PERIOD) self.mb_top5_err = ScalarMeter(cfg.LOG_PERIOD) @@ -500,6 +525,7 @@ def __init__(self, epoch_iters, cfg): self.num_top5_mis = 0 self.num_samples = 0 self.output_dir = cfg.OUTPUT_DIR + self.multi_loss = None def reset(self): """ @@ -508,11 +534,14 @@ def reset(self): self.loss.reset() self.loss_total = 0.0 self.lr = None + self.grad_norm = None self.mb_top1_err.reset() self.mb_top5_err.reset() self.num_top1_mis = 0 self.num_top5_mis = 0 self.num_samples = 0 + if self.multi_loss is not None: + self.multi_loss.reset() def iter_tic(self): """ @@ -532,7 +561,9 @@ def data_toc(self): self.data_timer.pause() self.net_timer.reset() - def update_stats(self, top1_err, top5_err, loss, lr, mb_size): + def update_stats( + self, top1_err, top5_err, loss, lr, grad_norm, mb_size, multi_loss=None + ): """ Update the current stats. Args: @@ -541,9 +572,11 @@ def update_stats(self, top1_err, top5_err, loss, lr, mb_size): loss (float): loss value. lr (float): learning rate. mb_size (int): mini batch size. + multi_loss (list): a list of values for multi-tasking losses. """ self.loss.add_value(loss) self.lr = lr + self.grad_norm = grad_norm self.loss_total += loss * mb_size self.num_samples += mb_size @@ -554,6 +587,23 @@ def update_stats(self, top1_err, top5_err, loss, lr, mb_size): # Aggregate stats self.num_top1_mis += top1_err * mb_size self.num_top5_mis += top5_err * mb_size + if multi_loss: + if self.multi_loss is None: + self.multi_loss = ListMeter(len(multi_loss)) + self.multi_loss.add_value(multi_loss) + if ( + self._cfg.TRAIN.KILL_LOSS_EXPLOSION_FACTOR > 0.0 + and len(self.loss.deque) > 6 + ): + prev_loss = 0.0 + for i in range(2, 7): + prev_loss += self.loss.deque[len(self.loss.deque) - i] + if loss > self._cfg.TRAIN.KILL_LOSS_EXPLOSION_FACTOR * prev_loss / 5.0: + raise RuntimeError( + "ERROR: Got Loss explosion of {} {}".format( + loss, datetime.datetime.now() + ) + ) def log_iter_stats(self, cur_epoch, cur_iter): """ @@ -569,9 +619,7 @@ def log_iter_stats(self, cur_epoch, cur_iter): ) eta = str(datetime.timedelta(seconds=int(eta_sec))) stats = { - "_type": "train_iter_{}".format( - "ssl" if self._cfg.TASK == "ssl" else "" - ), + "_type": "train_iter_{}".format("ssl" if self._cfg.TASK == "ssl" else ""), "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH), "iter": "{}/{}".format(cur_iter + 1, self.epoch_iters), "dt": self.iter_timer.seconds(), @@ -580,11 +628,16 @@ def log_iter_stats(self, cur_epoch, cur_iter): "eta": eta, "loss": self.loss.get_win_median(), "lr": self.lr, + "grad_norm": self.grad_norm, "gpu_mem": "{:.2f}G".format(misc.gpu_mem_usage()), } if not self._cfg.DATA.MULTI_LABEL: stats["top1_err"] = self.mb_top1_err.get_win_median() stats["top5_err"] = self.mb_top5_err.get_win_median() + if self.multi_loss is not None: + loss_list = self.multi_loss.get_value() + for idx, loss in enumerate(loss_list): + stats["loss_" + str(idx)] = loss logging.log_json_stats(stats) def log_epoch_stats(self, cur_epoch): @@ -598,15 +651,14 @@ def log_epoch_stats(self, cur_epoch): ) eta = str(datetime.timedelta(seconds=int(eta_sec))) stats = { - "_type": "train_epoch{}".format( - "_ssl" if self._cfg.TASK == "ssl" else "" - ), + "_type": "train_epoch{}".format("_ssl" if self._cfg.TASK == "ssl" else ""), "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH), "dt": self.iter_timer.seconds(), "dt_data": self.data_timer.seconds(), "dt_net": self.net_timer.seconds(), "eta": eta, "lr": self.lr, + "grad_norm": self.grad_norm, "gpu_mem": "{:.2f}G".format(misc.gpu_mem_usage()), "RAM": "{:.2f}/{:.2f}G".format(*misc.cpu_mem_usage()), } @@ -617,10 +669,14 @@ def log_epoch_stats(self, cur_epoch): stats["top1_err"] = top1_err stats["top5_err"] = top5_err stats["loss"] = avg_loss + if self.multi_loss is not None: + avg_loss_list = self.multi_loss.get_global_avg() + for idx, loss in enumerate(avg_loss_list): + stats["loss_" + str(idx)] = loss logging.log_json_stats(stats, self.output_dir) -class ValMeter(object): +class ValMeter: """ Measures validation stats. """ @@ -720,9 +776,7 @@ def log_iter_stats(self, cur_epoch, cur_iter): eta_sec = self.iter_timer.seconds() * (self.max_iter - cur_iter - 1) eta = str(datetime.timedelta(seconds=int(eta_sec))) stats = { - "_type": "val_iter{}".format( - "_ssl" if self._cfg.TASK == "ssl" else "" - ), + "_type": "val_iter{}".format("_ssl" if self._cfg.TASK == "ssl" else ""), "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH), "iter": "{}/{}".format(cur_iter + 1, self.max_iter), "time_diff": self.iter_timer.seconds(), @@ -741,9 +795,7 @@ def log_epoch_stats(self, cur_epoch): cur_epoch (int): the number of current epoch. """ stats = { - "_type": "val_epoch{}".format( - "_ssl" if self._cfg.TASK == "ssl" else "" - ), + "_type": "val_epoch{}".format("_ssl" if self._cfg.TASK == "ssl" else ""), "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH), "time_diff": self.iter_timer.seconds(), "gpu_mem": "{:.2f}G".format(misc.gpu_mem_usage()), diff --git a/slowfast/utils/metrics.py b/slowfast/utils/metrics.py index c1303cc35..78532e5cc 100644 --- a/slowfast/utils/metrics.py +++ b/slowfast/utils/metrics.py @@ -22,9 +22,9 @@ def topks_correct(preds, labels, ks): topks_correct (list): list of numbers, where the `i`-th entry corresponds to the number of top-`ks[i]` correct predictions. """ - assert preds.size(0) == labels.size( - 0 - ), "Batch dim of predictions and labels must match" + assert preds.size(0) == labels.size(0), ( + "Batch dim of predictions and labels must match" + ) # Find the top max_k predictions for each sample _top_max_k_vals, top_max_k_inds = torch.topk( preds, max(ks), dim=1, largest=True, sorted=True diff --git a/slowfast/utils/misc.py b/slowfast/utils/misc.py index 5920546d4..b38f665d2 100644 --- a/slowfast/utils/misc.py +++ b/slowfast/utils/misc.py @@ -2,23 +2,24 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import json -import logging import math -import numpy as np import os from datetime import datetime + +import numpy as np import psutil +import slowfast.utils.logging as logging +import slowfast.utils.multiprocessing as mpu import torch +import torchvision.io as io from fvcore.nn.activation_count import activation_count from fvcore.nn.flop_count import flop_count from matplotlib import pyplot as plt -from torch import nn - -import slowfast.utils.logging as logging -import slowfast.utils.multiprocessing as mpu from slowfast.datasets.utils import pack_pathway_output from slowfast.models.batchnorm_helper import SubBatchNorm3d from slowfast.utils.env import pathmgr +from torch import nn +from torchvision.utils import make_grid logger = logging.get_logger(__name__) @@ -90,7 +91,7 @@ def _get_model_analysis_input(cfg, use_train_input): """ rgb_dimension = 3 if use_train_input: - if cfg.TRAIN.DATASET in ["imagenet", "imagenetprefetch"]: + if "imagenet" in cfg.TRAIN.DATASET: input_tensors = torch.rand( rgb_dimension, cfg.DATA.TRAIN_CROP_SIZE, @@ -104,7 +105,7 @@ def _get_model_analysis_input(cfg, use_train_input): cfg.DATA.TRAIN_CROP_SIZE, ) else: - if cfg.TEST.DATASET in ["imagenet", "imagenetprefetch"]: + if "imagenet" in cfg.TEST.DATASET: input_tensors = torch.rand( rgb_dimension, cfg.DATA.TEST_CROP_SIZE, @@ -181,13 +182,11 @@ def log_model_info(model, cfg, use_train_input=True): log info for testing. """ logger.info("Model:\n{}".format(model)) - logger.info("Params: {:,}".format(params_count(model))) + params = params_count(model) + logger.info("Params: {:,}".format(params)) logger.info("Mem: {:,} MB".format(gpu_mem_usage())) - logger.info( - "Flops: {:,} G".format( - get_model_stats(model, cfg, "flop", use_train_input) - ) - ) + flops = get_model_stats(model, cfg, "flop", use_train_input) + logger.info("Flops: {:,} G".format(flops)) logger.info( "Activations: {:,} M".format( get_model_stats(model, cfg, "activation", use_train_input) @@ -195,6 +194,7 @@ def log_model_info(model, cfg, use_train_input=True): ) logger.info("nvidia-smi") os.system("nvidia-smi") + return flops, params def is_eval_epoch(cfg, cur_epoch, multigrid_schedule): @@ -212,9 +212,7 @@ def is_eval_epoch(cfg, cur_epoch, multigrid_schedule): prev_epoch = 0 for s in multigrid_schedule: if cur_epoch < s[-1]: - period = max( - (s[-1] - prev_epoch) // cfg.MULTIGRID.EVAL_FREQ + 1, 1 - ) + period = max((s[-1] - prev_epoch) // cfg.MULTIGRID.EVAL_FREQ + 1, 1) return (s[-1] - 1 - cur_epoch) % period == 0 prev_epoch = s[-1] @@ -251,6 +249,115 @@ def plot_input(tensor, bboxes=(), texts=(), path="./tmp_vis.png"): f.savefig(path) +def plot_input_normed( + tensor, + bboxes=(), + texts=(), + path="./tmp_vis.png", + folder_path="", + make_grids=False, + output_video=False, +): + """ + Plot the input tensor with the optional bounding box and save it to disk. + Args: + tensor (tensor): a tensor with shape of `NxCxHxW`. + bboxes (tuple): bounding boxes with format of [[x, y, h, w]]. + texts (tuple): a tuple of string to plot. + path (str): path to the image to save to. + """ + tensor = tensor.float() + try: + os.mkdir(folder_path) + except Exception as e: + pass + tensor = convert_normalized_images(tensor) + if output_video: + # assert make_grids, "video needs to have make_grids on" + assert tensor.ndim == 5 + sz = tensor.shape + + if make_grids: + vid = tensor.reshape([sz[0], sz[1] * sz[2], sz[3], sz[4]]) + vid = make_grid(vid, padding=8, pad_value=1.0, nrow=sz[0]) + vid = vid.reshape([sz[1], sz[2], vid.shape[1], vid.shape[2]]) + else: + vid = tensor.reshape([sz[0] * sz[1], sz[2], sz[3], sz[4]]) + + vid = vid.permute([0, 2, 3, 1]) + vid *= 255.0 + vid = vid.to(torch.uint8) + fps = 30.0 * vid.shape[0] / 64.0 + io.video.write_video(path, vid, fps, video_codec="libx264") + elif make_grids: + if tensor.ndim > 4 and tensor.shape[0] == 1: + tensor = tensor.squeeze() + nrow = 1 + elif tensor.ndim == 5: + nrow = tensor.shape[1] + tensor = tensor.reshape( + shape=(-1, tensor.shape[2], tensor.shape[3], tensor.shape[4]) + ) + vis2 = ( + make_grid(tensor, padding=8, pad_value=1.0, nrow=nrow) + .permute(1, 2, 0) + .cpu() + .numpy() + ) + plt.imsave(fname=path, arr=vis2, format="png") + else: + f, ax = plt.subplots( + nrows=tensor.shape[0], + ncols=tensor.shape[1], + figsize=(10 * tensor.shape[1], 10 * tensor.shape[0]), + ) + + if tensor.shape[0] == 1: + for i in range(tensor.shape[1]): + ax[i].axis("off") + ax[i].imshow(tensor[0][i].permute(1, 2, 0)) + # ax[1][0].axis('off') + if bboxes is not None and len(bboxes) > i: + for box in bboxes[i]: + x1, y1, x2, y2 = box + ax[i].vlines(x1, y1, y2, colors="g", linestyles="solid") + ax[i].vlines(x2, y1, y2, colors="g", linestyles="solid") + ax[i].hlines(y1, x1, x2, colors="g", linestyles="solid") + ax[i].hlines(y2, x1, x2, colors="g", linestyles="solid") + + if texts is not None and len(texts) > i: + ax[i].text(0, 0, texts[i]) + else: + for i in range(tensor.shape[0]): + for j in range(tensor.shape[1]): + ax[i][j].axis("off") + ax[i][j].imshow(tensor[i][j].permute(1, 2, 0)) + # ax[1][0].axis('off') + if bboxes is not None and len(bboxes) > i: + for box in bboxes[i]: + x1, y1, x2, y2 = box + ax[i].vlines(x1, y1, y2, colors="g", linestyles="solid") + ax[i].vlines(x2, y1, y2, colors="g", linestyles="solid") + ax[i].hlines(y1, x1, x2, colors="g", linestyles="solid") + ax[i].hlines(y2, x1, x2, colors="g", linestyles="solid") + + if texts is not None and len(texts) > i: + ax[i].text(0, 0, texts[i]) + print(f"{path}") + f.tight_layout(pad=0.0) + with pathmgr.open(path, "wb") as h: + f.savefig(h) + + +def convert_normalized_images(tensor): + tensor = tensor * 0.225 + tensor = tensor + 0.45 + + tensor = tensor.clamp(min=0.0, max=1.0) + + return tensor + + def frozen_bn_stats(model): """ Set all the bn layers to eval mode. @@ -349,17 +456,11 @@ def get_class_names(path, parent_path=None, subset_path=None): with pathmgr.open(parent_path, "r") as f: d_parent = json.load(f) except EnvironmentError as err: - print( - "Fail to load file from {} with error {}".format( - parent_path, err - ) - ) + print("Fail to load file from {} with error {}".format(parent_path, err)) return class_parent = {} for parent, children in d_parent.items(): - indices = [ - class2idx[c] for c in children if class2idx.get(c) is not None - ] + indices = [class2idx[c] for c in children if class2idx.get(c) is not None] class_parent[parent] = indices subset_ids = None @@ -373,11 +474,7 @@ def get_class_names(path, parent_path=None, subset_path=None): if class2idx.get(name) is not None ] except EnvironmentError as err: - print( - "Fail to load file from {} with error {}".format( - subset_path, err - ) - ) + print("Fail to load file from {} with error {}".format(subset_path, err)) return return class_names, class_parent, subset_ids diff --git a/slowfast/utils/multigrid.py b/slowfast/utils/multigrid.py index 0cb0f8192..85f4050e1 100644 --- a/slowfast/utils/multigrid.py +++ b/slowfast/utils/multigrid.py @@ -4,13 +4,12 @@ """Helper functions for multigrid training.""" import numpy as np - import slowfast.utils.logging as logging logger = logging.get_logger(__name__) -class MultigridSchedule(object): +class MultigridSchedule: """ This class defines multigrid training schedule and update cfg accordingly. """ @@ -37,12 +36,8 @@ def init_multigrid(self, cfg): self.schedule = self.get_long_cycle_schedule(cfg) cfg.SOLVER.STEPS = [0] + [s[-1] for s in self.schedule] # Fine-tuning phase. - cfg.SOLVER.STEPS[-1] = ( - cfg.SOLVER.STEPS[-2] + cfg.SOLVER.STEPS[-1] - ) // 2 - cfg.SOLVER.LRS = [ - cfg.SOLVER.GAMMA ** s[0] * s[1][0] for s in self.schedule - ] + cfg.SOLVER.STEPS[-1] = (cfg.SOLVER.STEPS[-2] + cfg.SOLVER.STEPS[-1]) // 2 + cfg.SOLVER.LRS = [cfg.SOLVER.GAMMA ** s[0] * s[1][0] for s in self.schedule] # Fine-tuning phase. cfg.SOLVER.LRS = cfg.SOLVER.LRS[:-1] + [ cfg.SOLVER.LRS[-2], @@ -73,18 +68,14 @@ def update_long_cycle(self, cfg, cur_epoch): cfg (configs): the updated cfg. changed (bool): do we change long cycle shape at this epoch? """ - base_b, base_t, base_s = get_current_long_cycle_shape( - self.schedule, cur_epoch - ) + base_b, base_t, base_s = get_current_long_cycle_shape(self.schedule, cur_epoch) if base_s != cfg.DATA.TRAIN_CROP_SIZE or base_t != cfg.DATA.NUM_FRAMES: - cfg.DATA.NUM_FRAMES = base_t cfg.DATA.TRAIN_CROP_SIZE = base_s cfg.TRAIN.BATCH_SIZE = base_b * cfg.MULTIGRID.DEFAULT_B bs_factor = ( - float(cfg.TRAIN.BATCH_SIZE / cfg.NUM_GPUS) - / cfg.MULTIGRID.BN_BASE_SIZE + float(cfg.TRAIN.BATCH_SIZE / cfg.NUM_GPUS) / cfg.MULTIGRID.BN_BASE_SIZE ) if bs_factor < 1: @@ -102,9 +93,7 @@ def update_long_cycle(self, cfg, cur_epoch): logger.info("Long cycle updates:") logger.info("\tBN.NORM_TYPE: {}".format(cfg.BN.NORM_TYPE)) if cfg.BN.NORM_TYPE == "sync_batchnorm": - logger.info( - "\tBN.NUM_SYNC_DEVICES: {}".format(cfg.BN.NUM_SYNC_DEVICES) - ) + logger.info("\tBN.NUM_SYNC_DEVICES: {}".format(cfg.BN.NUM_SYNC_DEVICES)) elif cfg.BN.NORM_TYPE == "sub_batchnorm": logger.info("\tBN.NUM_SPLITS: {}".format(cfg.BN.NUM_SPLITS)) logger.info("\tTRAIN.BATCH_SIZE: {}".format(cfg.TRAIN.BATCH_SIZE)) @@ -113,9 +102,7 @@ def update_long_cycle(self, cfg, cur_epoch): cfg.DATA.NUM_FRAMES, cfg.MULTIGRID.LONG_CYCLE_SAMPLING_RATE ) ) - logger.info( - "\tDATA.TRAIN_CROP_SIZE: {}".format(cfg.DATA.TRAIN_CROP_SIZE) - ) + logger.info("\tDATA.TRAIN_CROP_SIZE: {}".format(cfg.DATA.TRAIN_CROP_SIZE)) return cfg, True else: return cfg, False @@ -134,9 +121,7 @@ def get_long_cycle_schedule(self, cfg): steps = cfg.SOLVER.STEPS - default_size = float( - cfg.DATA.NUM_FRAMES * cfg.DATA.TRAIN_CROP_SIZE**2 - ) + default_size = float(cfg.DATA.NUM_FRAMES * cfg.DATA.TRAIN_CROP_SIZE**2) default_iters = steps[-1] # Get shapes and average batch size for each long cycle shape. @@ -149,13 +134,11 @@ def get_long_cycle_schedule(self, cfg): shapes = [ [ base_t, - cfg.MULTIGRID.DEFAULT_S - * cfg.MULTIGRID.SHORT_CYCLE_FACTORS[0], + cfg.MULTIGRID.DEFAULT_S * cfg.MULTIGRID.SHORT_CYCLE_FACTORS[0], ], [ base_t, - cfg.MULTIGRID.DEFAULT_S - * cfg.MULTIGRID.SHORT_CYCLE_FACTORS[1], + cfg.MULTIGRID.DEFAULT_S * cfg.MULTIGRID.SHORT_CYCLE_FACTORS[1], ], [base_t, base_s], ] @@ -177,9 +160,7 @@ def get_long_cycle_schedule(self, cfg): step_epochs = steps[step_index + 1] - steps[step_index] for long_cycle_index, shapes in enumerate(all_shapes): - cur_epochs = ( - step_epochs * avg_bs[long_cycle_index] / sum(avg_bs) - ) + cur_epochs = step_epochs * avg_bs[long_cycle_index] / sum(avg_bs) cur_iters = cur_epochs / avg_bs[long_cycle_index] total_iters += cur_iters diff --git a/slowfast/visualization/async_predictor.py b/slowfast/visualization/async_predictor.py index bb11ef15b..8309e5dbf 100644 --- a/slowfast/visualization/async_predictor.py +++ b/slowfast/visualization/async_predictor.py @@ -2,12 +2,12 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import atexit -import numpy as np import queue -import torch -import torch.multiprocessing as mp +import numpy as np import slowfast.utils.logging as logging +import torch +import torch.multiprocessing as mp from slowfast.datasets import cv2_transform from slowfast.visualization.predictor import Predictor @@ -170,9 +170,7 @@ def __init__(self, video_vis, n_workers=None): self.put_id = -1 for _ in range(max(num_workers, 1)): self.procs.append( - AsyncVis._VisWorker( - video_vis, self.task_queue, self.result_queue - ) + AsyncVis._VisWorker(video_vis, self.task_queue, self.result_queue) ) for p in self.procs: @@ -246,9 +244,7 @@ def __init__(self, cfg, async_vis): slowfast/config/defaults.py async_vis (AsyncVis object): asynchronous visualizer. """ - self.model = AsycnActionPredictor( - cfg=cfg, result_queue=async_vis.task_queue - ) + self.model = AsycnActionPredictor(cfg=cfg, result_queue=async_vis.task_queue) self.async_vis = async_vis def put(self, task): diff --git a/slowfast/visualization/ava_demo_precomputed_boxes.py b/slowfast/visualization/ava_demo_precomputed_boxes.py index 9ec218a24..dd42fcece 100644 --- a/slowfast/visualization/ava_demo_precomputed_boxes.py +++ b/slowfast/visualization/ava_demo_precomputed_boxes.py @@ -1,14 +1,14 @@ #!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -import numpy as np import os -import cv2 -import torch -import tqdm +import cv2 +import numpy as np import slowfast.utils.checkpoint as cu import slowfast.utils.logging as logging +import torch +import tqdm from slowfast.datasets.ava_helper import parse_bboxes_file from slowfast.datasets.cv2_transform import scale, scale_boxes from slowfast.datasets.utils import get_sequence @@ -113,9 +113,7 @@ def get_input_clip(self, keyframe_idx): frame = scale(self.cfg.DATA.TEST_CROP_SIZE, frame) clip.append(frame) else: - logger.error( - "Unable to read frame. Duplicating previous frame." - ) + logger.error("Unable to read frame. Duplicating previous frame.") clip.append(clip[-1]) clip = process_cv2_inputs(clip, self.cfg) @@ -136,9 +134,7 @@ def get_predictions(self): # Print config. logger.info("Run demo with config:") logger.info(self.cfg) - assert ( - self.cfg.NUM_GPUS <= 1 - ), "Cannot run demo visualization on multiple GPUs." + assert self.cfg.NUM_GPUS <= 1, "Cannot run demo visualization on multiple GPUs." # Build the video model and print model statistics. model = build_model(self.cfg) @@ -149,9 +145,7 @@ def get_predictions(self): cu.load_test_checkpoint(self.cfg, model) logger.info("Finish loading model weights") logger.info("Start making predictions for precomputed boxes.") - for keyframe_idx, boxes_and_labels in tqdm.tqdm( - self.pred_boxes.items() - ): + for keyframe_idx, boxes_and_labels in tqdm.tqdm(self.pred_boxes.items()): inputs = self.get_input_clip(keyframe_idx) boxes = boxes_and_labels[0] boxes = torch.from_numpy(np.array(boxes)).float() @@ -221,8 +215,7 @@ def draw_video(self): ] draw_range_repeat = [ draw_range[0], - (draw_range[1] - draw_range[0]) * self.no_frames_repeat - + draw_range[0], + (draw_range[1] - draw_range[0]) * self.no_frames_repeat + draw_range[0], ] prev_buffer = [] prev_end_idx = 0 @@ -280,9 +273,7 @@ def draw_video(self): repeat = 1 current_draw_range = draw_range_repeat # Make sure draw range does not fall out of end of clip. - current_draw_range[1] = min( - current_draw_range[1], len(clip) - 1 - ) + current_draw_range[1] = min(current_draw_range[1], len(clip) - 1) ground_truth = boxes[0] bboxes = boxes[1] label = boxes[2] diff --git a/slowfast/visualization/demo_loader.py b/slowfast/visualization/demo_loader.py index b6e7ef4e9..c2fb7ed78 100644 --- a/slowfast/visualization/demo_loader.py +++ b/slowfast/visualization/demo_loader.py @@ -6,8 +6,8 @@ import queue import threading import time -import cv2 +import cv2 import slowfast.utils.logging as logging from slowfast.visualization.utils import TaskInfo @@ -25,14 +25,12 @@ def __init__(self, cfg): cfg (CfgNode): configs. Details can be found in slowfast/config/defaults.py """ - assert ( - cfg.DEMO.WEBCAM > -1 or cfg.DEMO.INPUT_VIDEO != "" - ), "Must specify a data source as input." - - self.source = ( - cfg.DEMO.WEBCAM if cfg.DEMO.WEBCAM > -1 else cfg.DEMO.INPUT_VIDEO + assert cfg.DEMO.WEBCAM > -1 or cfg.DEMO.INPUT_VIDEO != "", ( + "Must specify a data source as input." ) + self.source = cfg.DEMO.WEBCAM if cfg.DEMO.WEBCAM > -1 else cfg.DEMO.INPUT_VIDEO + self.display_width = cfg.DEMO.DISPLAY_WIDTH self.display_height = cfg.DEMO.DISPLAY_HEIGHT @@ -156,14 +154,12 @@ def __init__(self, cfg): cfg (CfgNode): configs. Details can be found in slowfast/config/defaults.py """ - assert ( - cfg.DEMO.WEBCAM > -1 or cfg.DEMO.INPUT_VIDEO != "" - ), "Must specify a data source as input." - - self.source = ( - cfg.DEMO.WEBCAM if cfg.DEMO.WEBCAM > -1 else cfg.DEMO.INPUT_VIDEO + assert cfg.DEMO.WEBCAM > -1 or cfg.DEMO.INPUT_VIDEO != "", ( + "Must specify a data source as input." ) + self.source = cfg.DEMO.WEBCAM if cfg.DEMO.WEBCAM > -1 else cfg.DEMO.INPUT_VIDEO + self.display_width = cfg.DEMO.DISPLAY_WIDTH self.display_height = cfg.DEMO.DISPLAY_HEIGHT @@ -251,9 +247,7 @@ def put_fn(self): self.buffer = frames[-self.buffer_size :] task.add_frames(self.put_id + 1, frames) - task.num_buffer_frames = ( - 0 if self.put_id == -1 else self.buffer_size - ) + task.num_buffer_frames = 0 if self.put_id == -1 else self.buffer_size with self.put_id_lock: self.put_id += 1 self.not_end = was_read diff --git a/slowfast/visualization/gradcam_utils.py b/slowfast/visualization/gradcam_utils.py index 07c0fa28f..120aab3c9 100644 --- a/slowfast/visualization/gradcam_utils.py +++ b/slowfast/visualization/gradcam_utils.py @@ -2,10 +2,9 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import matplotlib.pyplot as plt +import slowfast.datasets.utils as data_utils import torch import torch.nn.functional as F - -import slowfast.datasets.utils as data_utils from slowfast.visualization.utils import get_layer @@ -16,9 +15,7 @@ class GradCAM: https://arxiv.org/pdf/1610.02391.pdf """ - def __init__( - self, model, target_layers, data_mean, data_std, colormap="viridis" - ): + def __init__(self, model, target_layers, data_mean, data_std, colormap="viridis"): """ Args: model (model): the model to be used. @@ -78,9 +75,9 @@ def _calculate_localization_map(self, inputs, labels=None): each corresponding input. preds (tensor): shape (n_instances, n_class). Model predictions for `inputs`. """ - assert len(inputs) == len( - self.target_layers - ), "Must register the same number of target layers as the number of input pathways." + assert len(inputs) == len(self.target_layers), ( + "Must register the same number of target layers as the number of input pathways." + ) input_clone = [inp.clone() for inp in inputs] preds = self.model(input_clone) @@ -105,9 +102,7 @@ def _calculate_localization_map(self, inputs, labels=None): weights = torch.mean(gradients.view(B, C, Tg, -1), dim=3) weights = weights.view(B, C, Tg, 1, 1) - localization_map = torch.sum( - weights * activations, dim=1, keepdim=True - ) + localization_map = torch.sum(weights * activations, dim=1, keepdim=True) localization_map = F.relu(localization_map) localization_map = F.interpolate( localization_map, @@ -116,12 +111,8 @@ def _calculate_localization_map(self, inputs, labels=None): align_corners=False, ) localization_map_min, localization_map_max = ( - torch.min(localization_map.view(B, -1), dim=-1, keepdim=True)[ - 0 - ], - torch.max(localization_map.view(B, -1), dim=-1, keepdim=True)[ - 0 - ], + torch.min(localization_map.view(B, -1), dim=-1, keepdim=True)[0], + torch.max(localization_map.view(B, -1), dim=-1, keepdim=True)[0], ) localization_map_min = torch.reshape( localization_map_min, shape=(B, 1, 1, 1, 1) diff --git a/slowfast/visualization/prediction_vis.py b/slowfast/visualization/prediction_vis.py index e2189b78b..e466e1401 100644 --- a/slowfast/visualization/prediction_vis.py +++ b/slowfast/visualization/prediction_vis.py @@ -2,11 +2,10 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import numpy as np -import torch - import slowfast.datasets.utils as data_utils import slowfast.utils.logging as logging import slowfast.visualization.tensorboard_vis as tb +import torch from slowfast.utils.misc import get_class_names from slowfast.visualization.video_visualizer import VideoVisualizer @@ -100,9 +99,7 @@ def add_video(vid, preds, tag, true_class_name): vid = self.video_vis.draw_clip(vid, preds) vid = torch.from_numpy(np.array(vid)).permute(0, 3, 1, 2) vid = torch.unsqueeze(vid, dim=0) - self.writer.add_video( - vid, tag="{}: {}".format(tag, true_class_name) - ) + self.writer.add_video(vid, tag="{}: {}".format(tag, true_class_name)) mask = self._pick_wrong_preds(labels, preds) video_indices = torch.squeeze(mask.nonzero(), dim=-1) @@ -113,8 +110,7 @@ def add_video(vid, preds, tag, true_class_name): add_video( video_input[pathway][vid_idx], preds=preds[vid_idx], - tag=self.tag - + "/Video {}, Pathway {}".format(cur_vid_idx, pathway), + tag=self.tag + "/Video {}, Pathway {}".format(cur_vid_idx, pathway), true_class_name=self.class_names[labels[vid_idx]], ) diff --git a/slowfast/visualization/predictor.py b/slowfast/visualization/predictor.py index 3007aa58f..d5419290a 100644 --- a/slowfast/visualization/predictor.py +++ b/slowfast/visualization/predictor.py @@ -2,13 +2,13 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import queue + import cv2 +import slowfast.utils.checkpoint as cu import torch from detectron2 import model_zoo from detectron2.config import get_cfg from detectron2.engine import DefaultPredictor - -import slowfast.utils.checkpoint as cu from slowfast.datasets import cv2_transform from slowfast.models import build_model from slowfast.utils import logging @@ -30,9 +30,7 @@ def __init__(self, cfg, gpu_id=None): gpu_id (Optional[int]): GPU id. """ if cfg.NUM_GPUS: - self.gpu_id = ( - torch.cuda.current_device() if gpu_id is None else gpu_id - ) + self.gpu_id = torch.cuda.current_device() if gpu_id is None else gpu_id # Build the video model and print model statistics. self.model = build_model(cfg, gpu_id=gpu_id) @@ -69,13 +67,10 @@ def __call__(self, task): task.img_width, ) if self.cfg.DEMO.INPUT_FORMAT == "BGR": - frames = [ - cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames - ] + frames = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames] frames = [ - cv2_transform.scale(self.cfg.DATA.TEST_CROP_SIZE, frame) - for frame in frames + cv2_transform.scale(self.cfg.DATA.TEST_CROP_SIZE, frame) for frame in frames ] inputs = process_cv2_inputs(frames, self.cfg) if bboxes is not None: @@ -170,17 +165,13 @@ def __init__(self, cfg, gpu_id=None): """ self.cfg = get_cfg() - self.cfg.merge_from_file( - model_zoo.get_config_file(cfg.DEMO.DETECTRON2_CFG) - ) + self.cfg.merge_from_file(model_zoo.get_config_file(cfg.DEMO.DETECTRON2_CFG)) self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = cfg.DEMO.DETECTRON2_THRESH self.cfg.MODEL.WEIGHTS = cfg.DEMO.DETECTRON2_WEIGHTS self.cfg.INPUT.FORMAT = cfg.DEMO.INPUT_FORMAT if cfg.NUM_GPUS and gpu_id is None: gpu_id = torch.cuda.current_device() - self.cfg.MODEL.DEVICE = ( - "cuda:{}".format(gpu_id) if cfg.NUM_GPUS > 0 else "cpu" - ) + self.cfg.MODEL.DEVICE = "cuda:{}".format(gpu_id) if cfg.NUM_GPUS > 0 else "cpu" logger.info("Initialized Detectron2 Object Detection Model.") diff --git a/slowfast/visualization/tensorboard_vis.py b/slowfast/visualization/tensorboard_vis.py index 4f0b63aa8..044b66c69 100644 --- a/slowfast/visualization/tensorboard_vis.py +++ b/slowfast/visualization/tensorboard_vis.py @@ -4,20 +4,20 @@ import logging as log import math import os -import matplotlib.pyplot as plt -import torch -from torch.utils.tensorboard import SummaryWriter -from torchvision.utils import make_grid +import matplotlib.pyplot as plt import slowfast.utils.logging as logging import slowfast.visualization.utils as vis_utils +import torch from slowfast.utils.misc import get_class_names +from torch.utils.tensorboard import SummaryWriter +from torchvision.utils import make_grid logger = logging.get_logger(__name__) log.getLogger("matplotlib").setLevel(log.ERROR) -class TensorboardWriter(object): +class TensorboardWriter: """ Helper class to log information to Tensorboard. """ @@ -44,18 +44,14 @@ def __init__(self, cfg): self.hist_figsize = cfg.TENSORBOARD.HISTOGRAM.FIGSIZE if cfg.TENSORBOARD.LOG_DIR == "": - log_dir = os.path.join( - cfg.OUTPUT_DIR, "runs-{}".format(cfg.TRAIN.DATASET) - ) + log_dir = os.path.join(cfg.OUTPUT_DIR, "runs-{}".format(cfg.TRAIN.DATASET)) else: log_dir = os.path.join(cfg.OUTPUT_DIR, cfg.TENSORBOARD.LOG_DIR) self.writer = SummaryWriter(log_dir=log_dir) logger.info( "To see logged results in Tensorboard, please launch using the command \ - `tensorboard --port= --logdir {}`".format( - log_dir - ) + `tensorboard --port= --logdir {}`".format(log_dir) ) if cfg.TENSORBOARD.CLASS_NAMES_PATH != "": @@ -321,9 +317,7 @@ def plot_hist( figsize=figsize, ) writer.add_figure( - tag="Top {} predictions by classes/{}".format( - k, class_names[i] - ), + tag="Top {} predictions by classes/{}".format(k, class_names[i]), figure=hist, global_step=global_step, ) @@ -391,17 +385,14 @@ def add_ndim_array( reshaped_array = array.view(-1, *last2_dims) if heat_map: reshaped_array = [ - add_heatmap(array_2d).unsqueeze(0) - for array_2d in reshaped_array + add_heatmap(array_2d).unsqueeze(0) for array_2d in reshaped_array ] reshaped_array = torch.cat(reshaped_array, dim=0) else: reshaped_array = reshaped_array.unsqueeze(1) if nrow is None: nrow = int(math.sqrt(reshaped_array.size()[0])) - img_grid = make_grid( - reshaped_array, nrow, padding=1, normalize=normalize - ) + img_grid = make_grid(reshaped_array, nrow, padding=1, normalize=normalize) writer.add_image(name, img_grid, global_step=global_step) diff --git a/slowfast/visualization/utils.py b/slowfast/visualization/utils.py index a07be1303..b1334cbf2 100644 --- a/slowfast/visualization/utils.py +++ b/slowfast/visualization/utils.py @@ -2,12 +2,12 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import itertools -import numpy as np + import matplotlib.pyplot as plt +import numpy as np +import slowfast.utils.logging as logging import torch from sklearn.metrics import confusion_matrix - -import slowfast.utils.logging as logging from slowfast.datasets.utils import pack_pathway_output, tensor_normalize logger = logging.get_logger(__name__) @@ -246,9 +246,7 @@ def get_weights(self): if hasattr(cur_layer, "weight"): weights[layer] = cur_layer.weight.clone().detach() else: - logger.error( - "Layer {} does not have weight attribute.".format(layer) - ) + logger.error("Layer {} does not have weight attribute.".format(layer)) return weights diff --git a/slowfast/visualization/video_visualizer.py b/slowfast/visualization/video_visualizer.py index faa127294..1a70c6c9c 100644 --- a/slowfast/visualization/video_visualizer.py +++ b/slowfast/visualization/video_visualizer.py @@ -3,12 +3,12 @@ import itertools import logging as log -import numpy as np + import matplotlib.pyplot as plt +import numpy as np +import slowfast.utils.logging as logging import torch from detectron2.utils.visualizer import Visualizer - -import slowfast.utils.logging as logging from slowfast.utils.misc import get_class_names logger = logging.get_logger(__name__) @@ -36,9 +36,7 @@ def _create_text_labels(classes, scores, class_names, ground_truth=False): labels = ["[{}] {}".format("GT", label) for label in labels] elif scores is not None: assert len(classes) == len(scores) - labels = [ - "[{:.2f}] {}".format(s, label) for s, label in zip(scores, labels) - ] + labels = ["[{:.2f}] {}".format(s, label) for s, label in zip(scores, labels)] return labels @@ -135,9 +133,9 @@ def draw_multiple_text( """ if not isinstance(box_facecolors, list): box_facecolors = [box_facecolors] * len(text_ls) - assert len(box_facecolors) == len( - text_ls - ), "Number of colors provided is not equal to the number of text labels." + assert len(box_facecolors) == len(text_ls), ( + "Number of colors provided is not equal to the number of text labels." + ) if not font_size: font_size = self._default_font_size text_box_width = font_size + font_size // 2 @@ -211,9 +209,9 @@ def draw_multiple_text_upward( """ if not isinstance(box_facecolors, list): box_facecolors = [box_facecolors] * len(text_ls) - assert len(box_facecolors) == len( - text_ls - ), "Number of colors provided is not equal to the number of text labels." + assert len(box_facecolors) == len(text_ls), ( + "Number of colors provided is not equal to the number of text labels." + ) assert y_corner in [1, 3], "Y_corner must be either 1 or 3" if not font_size: @@ -264,9 +262,9 @@ def draw_multiple_text_downward( """ if not isinstance(box_facecolors, list): box_facecolors = [box_facecolors] * len(text_ls) - assert len(box_facecolors) == len( - text_ls - ), "Number of colors provided is not equal to the number of text labels." + assert len(box_facecolors) == len(text_ls), ( + "Number of colors provided is not equal to the number of text labels." + ) assert y_corner in [1, 3], "Y_corner must be either 1 or 3" if not font_size: @@ -378,9 +376,7 @@ def __init__( This is used for choosing predictions for visualization. """ - assert mode in ["top-k", "thres"], "Mode {} is not supported.".format( - mode - ) + assert mode in ["top-k", "thres"], "Mode {} is not supported.".format(mode) self.mode = mode self.num_classes = num_classes self.class_names, _, _ = get_class_names(class_names_path, None, None) @@ -465,15 +461,13 @@ def draw_one_frame( ) ) frame_visualizer = ImgVisualizer(frame, meta=None) - font_size = min( - max(np.sqrt(frame.shape[0] * frame.shape[1]) // 35, 5), 9 - ) + font_size = min(max(np.sqrt(frame.shape[0] * frame.shape[1]) // 35, 5), 9) top_corner = not ground_truth if bboxes is not None: - assert len(preds) == len( - bboxes - ), "Encounter {} predictions and {} bounding boxes".format( - len(preds), len(bboxes) + assert len(preds) == len(bboxes), ( + "Encounter {} predictions and {} bounding boxes".format( + len(preds), len(bboxes) + ) ) for i, box in enumerate(bboxes): text = text_labels[i] @@ -641,9 +635,9 @@ def _adjust_frames_type(self, frames): frames (list of frames): list of frames in range [0, 1]. adjusted (bool): whether the original frames need adjusted. """ - assert ( - frames is not None and len(frames) != 0 - ), "Frames does not contain any values" + assert frames is not None and len(frames) != 0, ( + "Frames does not contain any values" + ) frames = np.array(frames) assert np.array(frames).ndim == 4, "Frames must have 4 dimensions" adjusted = False @@ -670,8 +664,6 @@ def _get_thres_array(self, common_class_names=None): else: common_class_ids = list(range(self.num_classes)) - thres_array = np.full( - shape=(self.num_classes,), fill_value=self.lower_thres - ) + thres_array = np.full(shape=(self.num_classes,), fill_value=self.lower_thres) thres_array[common_class_ids] = self.thres self.thres = torch.from_numpy(thres_array) diff --git a/tools/benchmark.py b/tools/benchmark.py old mode 100755 new mode 100644 index acdc12a99..0d542cdf8 --- a/tools/benchmark.py +++ b/tools/benchmark.py @@ -16,9 +16,7 @@ def main(): args = parse_args() cfg = load_config(args) - launch_job( - cfg=cfg, init_method=args.init_method, func=benchmark_data_loading - ) + launch_job(cfg=cfg, init_method=args.init_method, func=benchmark_data_loading) if __name__ == "__main__": diff --git a/tools/demo_net.py b/tools/demo_net.py index a7e98ebde..e0244fecf 100644 --- a/tools/demo_net.py +++ b/tools/demo_net.py @@ -1,11 +1,11 @@ #!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -import numpy as np import time + +import numpy as np import torch import tqdm - from slowfast.utils import logging from slowfast.visualization.async_predictor import AsyncDemo, AsyncVis from slowfast.visualization.ava_demo_precomputed_boxes import ( @@ -37,9 +37,7 @@ def run_demo(cfg, frame_provider): logger.info("Run demo with config:") logger.info(cfg) common_classes = ( - cfg.DEMO.COMMON_CLASS_NAMES - if len(cfg.DEMO.LABEL_FILE_PATH) != 0 - else None + cfg.DEMO.COMMON_CLASS_NAMES if len(cfg.DEMO.LABEL_FILE_PATH) != 0 else None ) video_vis = VideoVisualizer( @@ -62,9 +60,9 @@ def run_demo(cfg, frame_provider): seq_len = cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE - assert ( - cfg.DEMO.BUFFER_SIZE <= seq_len // 2 - ), "Buffer size cannot be greater than half of sequence length." + assert cfg.DEMO.BUFFER_SIZE <= seq_len // 2, ( + "Buffer size cannot be greater than half of sequence length." + ) num_task = 0 # Start reading frames. frame_provider.start() diff --git a/tools/run_net.py b/tools/run_net.py index 847448f69..4ac8a01bd 100644 --- a/tools/run_net.py +++ b/tools/run_net.py @@ -2,14 +2,14 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. """Wrapper to train and test a video classification model.""" + from slowfast.config.defaults import assert_and_infer_cfg from slowfast.utils.misc import launch_job from slowfast.utils.parser import load_config, parse_args - -from demo_net import demo -from test_net import test -from train_net import train -from visualization import visualize +from vision.fair.slowfast.tools.demo_net import demo +from vision.fair.slowfast.tools.test_net import test +from vision.fair.slowfast.tools.train_net import train +from vision.fair.slowfast.tools.visualization import visualize def main(): @@ -20,6 +20,7 @@ def main(): print("config files: {}".format(args.cfg_files)) for path_to_config in args.cfg_files: cfg = load_config(args, path_to_config) + cfg = assert_and_infer_cfg(cfg) # Perform training. if cfg.TRAIN.ENABLE: @@ -27,12 +28,17 @@ def main(): # Perform multi-clip testing. if cfg.TEST.ENABLE: - launch_job(cfg=cfg, init_method=args.init_method, func=test) + if cfg.TEST.NUM_ENSEMBLE_VIEWS == -1: + num_view_list = [1, 3, 5, 7, 10] + for num_view in num_view_list: + cfg.TEST.NUM_ENSEMBLE_VIEWS = num_view + launch_job(cfg=cfg, init_method=args.init_method, func=test) + else: + launch_job(cfg=cfg, init_method=args.init_method, func=test) # Perform model visualization. if cfg.TENSORBOARD.ENABLE and ( - cfg.TENSORBOARD.MODEL_VIS.ENABLE - or cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE + cfg.TENSORBOARD.MODEL_VIS.ENABLE or cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE ): launch_job(cfg=cfg, init_method=args.init_method, func=visualize) diff --git a/tools/test_net.py b/tools/test_net.py index ddf53dc02..46ef50deb 100644 --- a/tools/test_net.py +++ b/tools/test_net.py @@ -3,16 +3,16 @@ """Multi-view test a video classification model.""" -import numpy as np import os import pickle -import torch +import numpy as np import slowfast.utils.checkpoint as cu import slowfast.utils.distributed as du import slowfast.utils.logging as logging import slowfast.utils.misc as misc import slowfast.visualization.tensorboard_vis as tb +import torch from slowfast.datasets import loader from slowfast.models import build_model from slowfast.utils.env import pathmgr @@ -46,10 +46,7 @@ def perform_test(test_loader, model, test_meter, cfg, writer=None): model.eval() test_meter.iter_tic() - for cur_iter, (inputs, labels, video_idx, time, meta) in enumerate( - test_loader - ): - + for cur_iter, (inputs, labels, video_idx, time, meta) in enumerate(test_loader): if cfg.NUM_GPUS: # Transfer the data to the current GPU device. if isinstance(inputs, (list,)): @@ -75,12 +72,8 @@ def perform_test(test_loader, model, test_meter, cfg, writer=None): metadata = meta["metadata"] preds = preds.detach().cpu() if cfg.NUM_GPUS else preds.detach() - ori_boxes = ( - ori_boxes.detach().cpu() if cfg.NUM_GPUS else ori_boxes.detach() - ) - metadata = ( - metadata.detach().cpu() if cfg.NUM_GPUS else metadata.detach() - ) + ori_boxes = ori_boxes.detach().cpu() if cfg.NUM_GPUS else ori_boxes.detach() + metadata = metadata.detach().cpu() if cfg.NUM_GPUS else metadata.detach() if cfg.NUM_GPUS > 1: preds = torch.cat(du.all_gather_unaligned(preds), dim=0) @@ -118,7 +111,6 @@ def perform_test(test_loader, model, test_meter, cfg, writer=None): else: # Perform the forward pass. preds = model(inputs) - # Gather all the predictions across all the devices to perform ensemble. if cfg.NUM_GPUS > 1: preds, labels, video_idx = du.all_gather([preds, labels, video_idx]) @@ -128,10 +120,10 @@ def perform_test(test_loader, model, test_meter, cfg, writer=None): video_idx = video_idx.cpu() test_meter.iter_toc() - # Update and log stats. - test_meter.update_stats( - preds.detach(), labels.detach(), video_idx.detach() - ) + + if not cfg.VIS_MASK.ENABLE: + # Update and log stats. + test_meter.update_stats(preds.detach(), labels.detach(), video_idx.detach()) test_meter.log_iter_stats(cur_iter) test_meter.iter_tic() @@ -153,9 +145,7 @@ def perform_test(test_loader, model, test_meter, cfg, writer=None): with pathmgr.open(save_path, "wb") as f: pickle.dump([all_preds, all_labels], f) - logger.info( - "Successfully saved prediction results to {}".format(save_path) - ) + logger.info("Successfully saved prediction results to {}".format(save_path)) test_meter.finalize_metrics() return test_meter @@ -177,83 +167,101 @@ def test(cfg): # Setup logging format. logging.setup_logging(cfg.OUTPUT_DIR) - # Print config. - logger.info("Test with config:") - logger.info(cfg) + if len(cfg.TEST.NUM_TEMPORAL_CLIPS) == 0: + cfg.TEST.NUM_TEMPORAL_CLIPS = [cfg.TEST.NUM_ENSEMBLE_VIEWS] + + test_meters = [] + for num_view in cfg.TEST.NUM_TEMPORAL_CLIPS: + cfg.TEST.NUM_ENSEMBLE_VIEWS = num_view + + # Print config. + logger.info("Test with config:") + logger.info(cfg) + + # Build the video model and print model statistics. + model = build_model(cfg) + flops, params = 0.0, 0.0 + if du.is_master_proc() and cfg.LOG_MODEL_INFO: + model.eval() + flops, params = misc.log_model_info(model, cfg, use_train_input=False) + + if du.is_master_proc() and cfg.LOG_MODEL_INFO: + misc.log_model_info(model, cfg, use_train_input=False) + if ( + cfg.TASK == "ssl" + and cfg.MODEL.MODEL_NAME == "ContrastiveModel" + and cfg.CONTRASTIVE.KNN_ON + ): + train_loader = loader.construct_loader(cfg, "train") + if hasattr(model, "module"): + model.module.init_knn_labels(train_loader) + else: + model.init_knn_labels(train_loader) - # Build the video model and print model statistics. - model = build_model(cfg) + cu.load_test_checkpoint(cfg, model) - out_str_prefix = "lin" if cfg.MODEL.DETACH_FINAL_FC else "" + # Create video testing loaders. + test_loader = loader.construct_loader(cfg, "test") + logger.info("Testing model for {} iterations".format(len(test_loader))) - if du.is_master_proc() and cfg.LOG_MODEL_INFO: - misc.log_model_info(model, cfg, use_train_input=False) + if cfg.DETECTION.ENABLE: + assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE or cfg.NUM_GPUS == 0 + test_meter = AVAMeter(len(test_loader), cfg, mode="test") + else: + assert ( + test_loader.dataset.num_videos + % (cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS) + == 0 + ) + # Create meters for multi-view testing. + test_meter = TestMeter( + test_loader.dataset.num_videos + // (cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS), + cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS, + ( + cfg.MODEL.NUM_CLASSES + if not cfg.TASK == "ssl" + else cfg.CONTRASTIVE.NUM_CLASSES_DOWNSTREAM + ), + len(test_loader), + cfg.DATA.MULTI_LABEL, + cfg.DATA.ENSEMBLE_METHOD, + ) - if ( - cfg.TASK == "ssl" - and cfg.MODEL.MODEL_NAME == "ContrastiveModel" - and cfg.CONTRASTIVE.KNN_ON - ): - train_loader = loader.construct_loader(cfg, "train") - out_str_prefix = "knn" - if hasattr(model, "module"): - model.module.init_knn_labels(train_loader) + # Set up writer for logging to Tensorboard format. + if cfg.TENSORBOARD.ENABLE and du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS): + writer = tb.TensorboardWriter(cfg) else: - model.init_knn_labels(train_loader) - - cu.load_test_checkpoint(cfg, model) - - # Create video testing loaders. - test_loader = loader.construct_loader(cfg, "test") - logger.info("Testing model for {} iterations".format(len(test_loader))) - - if cfg.DETECTION.ENABLE: - assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE or cfg.NUM_GPUS == 0 - test_meter = AVAMeter(len(test_loader), cfg, mode="test") - else: - assert ( - test_loader.dataset.num_videos - % (cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS) - == 0 - ) - # Create meters for multi-view testing. - test_meter = TestMeter( - test_loader.dataset.num_videos - // (cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS), - cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS, - cfg.MODEL.NUM_CLASSES - if not cfg.TASK == "ssl" - else cfg.CONTRASTIVE.NUM_CLASSES_DOWNSTREAM, - len(test_loader), - cfg.DATA.MULTI_LABEL, - cfg.DATA.ENSEMBLE_METHOD, - ) + writer = None - # Set up writer for logging to Tensorboard format. - if cfg.TENSORBOARD.ENABLE and du.is_master_proc( - cfg.NUM_GPUS * cfg.NUM_SHARDS - ): - writer = tb.TensorboardWriter(cfg) - else: - writer = None - - # # Perform multi-view test on the entire dataset. - test_meter = perform_test(test_loader, model, test_meter, cfg, writer) - if writer is not None: - writer.close() - result_string = ( - "_a{}{}{} Top1 Acc: {} Top5 Acc: {} MEM: {:.2f} dataset: {}{}" - "".format( - out_str_prefix, - cfg.TEST.DATASET[0], - test_meter.stats["top1_acc"], - test_meter.stats["top1_acc"], - test_meter.stats["top5_acc"], - misc.gpu_mem_usage(), - cfg.TEST.DATASET[0], - cfg.MODEL.NUM_CLASSES, + # # Perform multi-view test on the entire dataset. + test_meter = perform_test(test_loader, model, test_meter, cfg, writer) + test_meters.append(test_meter) + if writer is not None: + writer.close() + + result_string_views = "_p{:.2f}_f{:.2f}".format(params / 1e6, flops) + + for view, test_meter in zip(cfg.TEST.NUM_TEMPORAL_CLIPS, test_meters): + logger.info( + "Finalized testing with {} temporal clips and {} spatial crops".format( + view, cfg.TEST.NUM_SPATIAL_CROPS + ) + ) + result_string_views += "_{}a{}".format(view, test_meter.stats["top1_acc"]) + + result_string = ( + "_p{:.2f}_f{:.2f}_{}a{} Top5 Acc: {} MEM: {:.2f} f: {:.4f}".format( + params / 1e6, + flops, + view, + test_meter.stats["top1_acc"], + test_meter.stats["top5_acc"], + misc.gpu_mem_usage(), + flops, + ) ) - ) - logger.info("testing done: {}".format(result_string)) - return result_string + logger.info("{}".format(result_string)) + logger.info("{}".format(result_string_views)) + return result_string + " \n " + result_string_views diff --git a/tools/train_net.py b/tools/train_net.py index cfc75aff4..4da44f7cc 100644 --- a/tools/train_net.py +++ b/tools/train_net.py @@ -4,11 +4,9 @@ """Train a video classification model.""" import math -import numpy as np import pprint -import torch -from fvcore.nn.precise_bn import get_bn_modules, update_bn_stats +import numpy as np import slowfast.models.losses as losses import slowfast.models.optimizer as optim import slowfast.utils.checkpoint as cu @@ -17,10 +15,15 @@ import slowfast.utils.metrics as metrics import slowfast.utils.misc as misc import slowfast.visualization.tensorboard_vis as tb +import torch +from fvcore.nn.precise_bn import get_bn_modules, update_bn_stats from slowfast.datasets import loader from slowfast.datasets.mixup import MixUp from slowfast.models import build_model -from slowfast.models.contrastive import cancel_swav_gradients +from slowfast.models.contrastive import ( + contrastive_forward, + contrastive_parameter_surgery, +) from slowfast.utils.meters import AVAMeter, EpochTimer, TrainMeter, ValMeter from slowfast.utils.multigrid import MultigridSchedule @@ -66,26 +69,12 @@ def train_epoch( num_classes=cfg.MODEL.NUM_CLASSES, ) - iters_noupdate = 0 - if ( - cfg.MODEL.MODEL_NAME == "ContrastiveModel" - and cfg.CONTRASTIVE.TYPE == "moco" - ): - assert ( - cfg.CONTRASTIVE.QUEUE_LEN % (cfg.TRAIN.BATCH_SIZE * cfg.NUM_SHARDS) - == 0 - ) - iters_noupdate = ( - cfg.CONTRASTIVE.QUEUE_LEN // cfg.TRAIN.BATCH_SIZE // cfg.NUM_SHARDS - ) if cfg.MODEL.FROZEN_BN: misc.frozen_bn_stats(model) # Explicitly declare reduction to mean. loss_fun = losses.get_loss_func(cfg.MODEL.LOSS_FUNC)(reduction="mean") - for cur_iter, (inputs, labels, index, time, meta) in enumerate( - train_loader - ): + for cur_iter, (inputs, labels, index, time, meta) in enumerate(train_loader): # Transfer the data to the current GPU device. if cfg.NUM_GPUS: if isinstance(inputs, (list,)): @@ -97,19 +86,19 @@ def train_epoch( inputs[i] = inputs[i].cuda(non_blocking=True) else: inputs = inputs.cuda(non_blocking=True) - labels = labels.cuda() + if not isinstance(labels, list): + labels = labels.cuda(non_blocking=True) + index = index.cuda(non_blocking=True) + time = time.cuda(non_blocking=True) for key, val in meta.items(): if isinstance(val, (list,)): for i in range(len(val)): val[i] = val[i].cuda(non_blocking=True) else: meta[key] = val.cuda(non_blocking=True) - index = index.cuda() - time = time.cuda() + batch_size = ( - inputs[0][0].size(0) - if isinstance(inputs[0], list) - else inputs[0].size(0) + inputs[0][0].size(0) if isinstance(inputs[0], list) else inputs[0].size(0) ) # Update the learning rate. epoch_exact = cur_epoch + float(cur_iter) / data_size @@ -122,7 +111,6 @@ def train_epoch( inputs[0] = samples with torch.cuda.amp.autocast(enabled=cfg.TRAIN.MIXED_PRECISION): - # Explicitly declare reduction to mean. perform_backward = True optimizer.zero_grad() @@ -139,43 +127,47 @@ def train_epoch( elif cfg.DETECTION.ENABLE: # Compute the predictions. preds = model(inputs, meta["boxes"]) + elif cfg.MASK.ENABLE: + preds, labels = model(inputs) else: preds = model(inputs) - if cfg.TASK == "ssl" and cfg.MODEL.MODEL_NAME == "ContrastiveModel": - labels = torch.zeros( - preds.size(0), dtype=labels.dtype, device=labels.device - ) + if cfg.TASK == "ssl" and cfg.MODEL.MODEL_NAME == "ContrastiveModel": + labels = torch.zeros( + preds.size(0), dtype=labels.dtype, device=labels.device + ) - if cfg.MODEL.MODEL_NAME == "ContrastiveModel" and partial_loss: - loss = partial_loss - else: - # Compute the loss. - loss = loss_fun(preds, labels) + if cfg.MODEL.MODEL_NAME == "ContrastiveModel" and partial_loss: + loss = partial_loss + else: + # Compute the loss. + loss = loss_fun(preds, labels) + + loss_extra = None + if isinstance(loss, (list, tuple)): + loss, loss_extra = loss # check Nan Loss. misc.check_nan_losses(loss) - if perform_backward: scaler.scale(loss).backward() # Unscales the gradients of optimizer's assigned params in-place scaler.unscale_(optimizer) # Clip gradients if necessary if cfg.SOLVER.CLIP_GRAD_VAL: - torch.nn.utils.clip_grad_value_( + grad_norm = torch.nn.utils.clip_grad_value_( model.parameters(), cfg.SOLVER.CLIP_GRAD_VAL ) elif cfg.SOLVER.CLIP_GRAD_L2NORM: - torch.nn.utils.clip_grad_norm_( + grad_norm = torch.nn.utils.clip_grad_norm_( model.parameters(), cfg.SOLVER.CLIP_GRAD_L2NORM ) - - model = cancel_swav_gradients(model, cfg, epoch_exact) - if cur_iter < iters_noupdate and cur_epoch == 0: # for e.g. MoCo - logger.info( - "Not updating parameters {}/{}".format(cur_iter, iters_noupdate) - ) else: - # Update the parameters. + grad_norm = optim.get_grad_norm_(model.parameters()) + # Update the parameters. (defaults to True) + model, update_param = contrastive_parameter_surgery( + model, cfg, epoch_exact, cur_iter + ) + if update_param: scaler.step(optimizer) scaler.update() @@ -209,8 +201,25 @@ def train_epoch( if cfg.DATA.MULTI_LABEL: # Gather all the predictions across all the devices. if cfg.NUM_GPUS > 1: - [loss] = du.all_reduce([loss]) - loss = loss.item() + loss, grad_norm = du.all_reduce([loss, grad_norm]) + loss, grad_norm = ( + loss.item(), + grad_norm.item(), + ) + elif cfg.MASK.ENABLE: + # Gather all the predictions across all the devices. + if cfg.NUM_GPUS > 1: + loss, grad_norm = du.all_reduce([loss, grad_norm]) + if loss_extra: + loss_extra = du.all_reduce(loss_extra) + loss, grad_norm, top1_err, top5_err = ( + loss.item(), + grad_norm.item(), + 0.0, + 0.0, + ) + if loss_extra: + loss_extra = [one_loss.item() for one_loss in loss_extra] else: # Compute the errors. num_topks_correct = metrics.topks_correct(preds, labels, (1, 5)) @@ -219,13 +228,14 @@ def train_epoch( ] # Gather all the predictions across all the devices. if cfg.NUM_GPUS > 1: - loss, top1_err, top5_err = du.all_reduce( - [loss.detach(), top1_err, top5_err] + loss, grad_norm, top1_err, top5_err = du.all_reduce( + [loss.detach(), grad_norm, top1_err, top5_err] ) # Copy the stats from GPU to CPU (sync point). - loss, top1_err, top5_err = ( + loss, grad_norm, top1_err, top5_err = ( loss.item(), + grad_norm.item(), top1_err.item(), top5_err.item(), ) @@ -236,10 +246,12 @@ def train_epoch( top5_err, loss, lr, + grad_norm, batch_size * max( cfg.NUM_GPUS, 1 ), # If running on CPU (cfg.NUM_GPUS == 1), use 1 to represent 1 CPU. + loss_extra, ) # write to tensorboard format if available. if writer is not None: @@ -252,22 +264,22 @@ def train_epoch( }, global_step=data_size * cur_epoch + cur_iter, ) - - torch.cuda.synchronize() train_meter.iter_toc() # do measure allreduce for this meter train_meter.log_iter_stats(cur_epoch, cur_iter) torch.cuda.synchronize() train_meter.iter_tic() del inputs + + # in case of fragmented memory + torch.cuda.empty_cache() + # Log epoch stats. train_meter.log_epoch_stats(cur_epoch) train_meter.reset() @torch.no_grad() -def eval_epoch( - val_loader, model, val_meter, cur_epoch, cfg, train_loader, writer -): +def eval_epoch(val_loader, model, val_meter, cur_epoch, cfg, train_loader, writer): """ Evaluate the model on the val set. Args: @@ -303,9 +315,7 @@ def eval_epoch( index = index.cuda() time = time.cuda() batch_size = ( - inputs[0][0].size(0) - if isinstance(inputs[0], list) - else inputs[0].size(0) + inputs[0][0].size(0) if isinstance(inputs[0], list) else inputs[0].size(0) ) val_meter.data_toc() @@ -340,9 +350,7 @@ def eval_epoch( ) yd, yi = model(inputs, index, time) K = yi.shape[1] - C = ( - cfg.CONTRASTIVE.NUM_CLASSES_DOWNSTREAM - ) # eg 400 for Kinetics400 + C = cfg.CONTRASTIVE.NUM_CLASSES_DOWNSTREAM # eg 400 for Kinetics400 candidates = train_labels.view(1, -1).expand(batch_size, -1) retrieval = torch.gather(candidates, 1, yi) retrieval_one_hot = torch.zeros((batch_size * K, C)).cuda() @@ -360,6 +368,8 @@ def eval_epoch( if cfg.NUM_GPUS > 1: preds, labels = du.all_gather([preds, labels]) else: + if cfg.DATA.IN22k_VAL_IN1K != "": + preds = preds[:, :1000] # Compute the errors. num_topks_correct = metrics.topks_correct(preds, labels, (1, 5)) @@ -400,74 +410,18 @@ def eval_epoch( # write to tensorboard format if available. if writer is not None: if cfg.DETECTION.ENABLE: - writer.add_scalars( - {"Val/mAP": val_meter.full_map}, global_step=cur_epoch - ) + writer.add_scalars({"Val/mAP": val_meter.full_map}, global_step=cur_epoch) else: all_preds = [pred.clone().detach() for pred in val_meter.all_preds] - all_labels = [ - label.clone().detach() for label in val_meter.all_labels - ] + all_labels = [label.clone().detach() for label in val_meter.all_labels] if cfg.NUM_GPUS: all_preds = [pred.cpu() for pred in all_preds] all_labels = [label.cpu() for label in all_labels] - writer.plot_eval( - preds=all_preds, labels=all_labels, global_step=cur_epoch - ) + writer.plot_eval(preds=all_preds, labels=all_labels, global_step=cur_epoch) val_meter.reset() -def contrastive_forward(model, cfg, inputs, index, time, epoch_exact, scaler): - if cfg.CONTRASTIVE.SEQUENTIAL: - perform_backward = False - mdl = model.module if hasattr(model, "module") else model - keys = ( - mdl.compute_key_feat( - inputs, - compute_predictor_keys=False, - batched_inference=True if len(inputs) < 2 else False, - ) - if cfg.CONTRASTIVE.TYPE == "moco" or cfg.CONTRASTIVE.TYPE == "byol" - else [None] * len(inputs) - ) - for k, vid in enumerate(inputs): - other_keys = keys[:k] + keys[k + 1 :] - time_cur = torch.cat( - [ - time[:, k : k + 1, :], - time[:, :k, :], - time[:, k + 1 :, :], - ], - 1, - ) # q, kpre, kpost - vids = [vid] - if ( - cfg.CONTRASTIVE.TYPE == "swav" - or cfg.CONTRASTIVE.TYPE == "simclr" - ): - if k < len(inputs) - 1: - vids = inputs[k : k + 2] - else: - break - lgt_k, loss_k = model( - vids, index, time_cur, epoch_exact, keys=other_keys - ) - scaler.scale(loss_k).backward() - if k == 0: - preds, partial_loss = lgt_k, loss_k.detach() - else: - preds = torch.cat([preds, lgt_k], dim=0) - partial_loss += loss_k.detach() - partial_loss /= len(inputs) * 2.0 # to have same loss as symm model - if cfg.CONTRASTIVE.TYPE == "moco": - mdl._dequeue_and_enqueue(keys) - else: - perform_backward = True - preds, partial_loss = model(inputs, index, time, epoch_exact, keys=None) - return model, preds, partial_loss, perform_backward - - def calculate_and_update_precise_bn(loader, model, num_iters=200, use_gpu=True): """ Update the stats in bn layers by calculate the precise stats. @@ -512,7 +466,7 @@ def build_trainer(cfg): # Build the video model and print model statistics. model = build_model(cfg) if du.is_master_proc() and cfg.LOG_MODEL_INFO: - misc.log_model_info(model, cfg, use_train_input=True) + flops, params = misc.log_model_info(model, cfg, use_train_input=True) # Construct the optimizer. optimizer = optim.construct_optimizer(model, cfg) @@ -520,9 +474,7 @@ def build_trainer(cfg): # Create the video train and val loaders. train_loader = loader.construct_loader(cfg, "train") val_loader = loader.construct_loader(cfg, "val") - precise_bn_loader = loader.construct_loader( - cfg, "train", is_precise_bn=True - ) + precise_bn_loader = loader.construct_loader(cfg, "train", is_precise_bn=True) # Create meters. train_meter = TrainMeter(len(train_loader), cfg) val_meter = ValMeter(len(val_loader), cfg) @@ -567,8 +519,9 @@ def train(cfg): # Build the video model and print model statistics. model = build_model(cfg) + flops, params = 0.0, 0.0 if du.is_master_proc() and cfg.LOG_MODEL_INFO: - misc.log_model_info(model, cfg, use_train_input=True) + flops, params = misc.log_model_info(model, cfg, use_train_input=True) # Construct the optimizer. optimizer = optim.construct_optimizer(model, cfg) @@ -614,6 +567,7 @@ def train(cfg): convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2", epoch_reset=cfg.TRAIN.CHECKPOINT_EPOCH_RESET, clear_name_pattern=cfg.TRAIN.CHECKPOINT_CLEAR_NAME_PATTERN, + image_init=cfg.TRAIN.CHECKPOINT_IN_INIT, ) start_epoch = checkpoint_epoch + 1 else: @@ -647,9 +601,7 @@ def train(cfg): val_meter = ValMeter(len(val_loader), cfg) # set up writer for logging to Tensorboard format. - if cfg.TENSORBOARD.ENABLE and du.is_master_proc( - cfg.NUM_GPUS * cfg.NUM_SHARDS - ): + if cfg.TENSORBOARD.ENABLE and du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS): writer = tb.TensorboardWriter(cfg) else: writer = None @@ -659,7 +611,6 @@ def train(cfg): epoch_timer = EpochTimer() for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH): - if cur_epoch > 0 and cfg.DATA.LOADER_CHUNK_SIZE > 0: num_chunks = math.ceil( cfg.DATA.LOADER_CHUNK_OVERALL_SIZE / cfg.DATA.LOADER_CHUNK_SIZE @@ -695,9 +646,7 @@ def train(cfg): else: last_checkpoint = cfg.TRAIN.CHECKPOINT_FILE_PATH logger.info("Load from {}".format(last_checkpoint)) - cu.load_checkpoint( - last_checkpoint, model, cfg.NUM_GPUS > 1, optimizer - ) + cu.load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1, optimizer) # Shuffle the dataset. loader.shuffle_dataset(train_loader, cur_epoch) @@ -724,9 +673,9 @@ def train(cfg): ) logger.info( f"For epoch {cur_epoch}, each iteraction takes " - f"{epoch_timer.last_epoch_time()/len(train_loader):.2f}s in average. " + f"{epoch_timer.last_epoch_time() / len(train_loader):.2f}s in average. " f"From epoch {start_epoch} to {cur_epoch}, each iteraction takes " - f"{epoch_timer.avg_epoch_time()/len(train_loader):.2f}s in average." + f"{epoch_timer.avg_epoch_time() / len(train_loader):.2f}s in average." ) is_checkp_epoch = ( @@ -737,8 +686,13 @@ def train(cfg): ) or cur_epoch == cfg.SOLVER.MAX_EPOCH - 1 ) - is_eval_epoch = misc.is_eval_epoch( - cfg, cur_epoch, None if multigrid is None else multigrid.schedule + is_eval_epoch = ( + misc.is_eval_epoch( + cfg, + cur_epoch, + None if multigrid is None else multigrid.schedule, + ) + and not cfg.MASK.ENABLE ) # Compute precise BN stats. @@ -776,12 +730,28 @@ def train(cfg): train_loader, writer, ) + if ( + start_epoch == cfg.SOLVER.MAX_EPOCH and not cfg.MASK.ENABLE + ): # final checkpoint load + eval_epoch(val_loader, model, val_meter, start_epoch, cfg, train_loader, writer) if writer is not None: writer.close() - result_string = "Top1 Acc: {:.2f} Top5 Acc: {:.2f} MEM: {:.2f}" "".format( - 100 - val_meter.min_top1_err, - 100 - val_meter.min_top5_err, - misc.gpu_mem_usage(), + result_string = ( + "_p{:.2f}_f{:.2f} _t{:.2f}_m{:.2f} _a{:.2f} Top5 Acc: {:.2f} MEM: {:.2f} f: {:.4f}" + "".format( + params / 1e6, + flops, + ( + epoch_timer.median_epoch_time() / 60.0 + if len(epoch_timer.epoch_times) + else 0.0 + ), + misc.gpu_mem_usage(), + 100 - val_meter.min_top1_err, + 100 - val_meter.min_top5_err, + misc.gpu_mem_usage(), + flops, + ) ) logger.info("training done: {}".format(result_string)) diff --git a/tools/visualization.py b/tools/visualization.py index 97449ec2f..35ecf398d 100644 --- a/tools/visualization.py +++ b/tools/visualization.py @@ -1,17 +1,17 @@ #!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -import numpy as np import pickle -import torch -import tqdm +import numpy as np import slowfast.datasets.utils as data_utils import slowfast.utils.checkpoint as cu import slowfast.utils.distributed as du import slowfast.utils.logging as logging import slowfast.utils.misc as misc import slowfast.visualization.tensorboard_vis as tb +import torch +import tqdm from slowfast.datasets import loader from slowfast.models import build_model from slowfast.utils.env import pathmgr @@ -62,8 +62,7 @@ def run_visualization(vis_loader, model, cfg, writer=None): ) if n_devices > 1: grad_cam_layer_ls = [ - "module/" + layer - for layer in cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST + "module/" + layer for layer in cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST ] else: grad_cam_layer_ls = cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST @@ -95,9 +94,7 @@ def run_visualization(vis_loader, model, cfg, writer=None): meta[key] = val.cuda(non_blocking=True) if cfg.DETECTION.ENABLE: - activations, preds = model_vis.get_activations( - inputs, meta["boxes"] - ) + activations, preds = model_vis.get_activations(inputs, meta["boxes"]) else: activations, preds = model_vis.get_activations(inputs) if cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.ENABLE: @@ -141,9 +138,7 @@ def run_visualization(vis_loader, model, cfg, writer=None): ): for path_idx, input_pathway in enumerate(cur_input): if cfg.TEST.DATASET == "ava" and cfg.AVA.BGR: - video = input_pathway[ - cur_batch_idx, [2, 1, 0], ... - ] + video = input_pathway[cur_batch_idx, [2, 1, 0], ...] else: video = input_pathway[cur_batch_idx] @@ -156,9 +151,7 @@ def run_visualization(vis_loader, model, cfg, writer=None): else: # Permute from (T, C, H, W) to (T, H, W, C) video = video.permute(0, 2, 3, 1) - bboxes = ( - None if cur_boxes is None else cur_boxes[:, 1:] - ) + bboxes = None if cur_boxes is None else cur_boxes[:, 1:] cur_prediction = ( cur_preds if cfg.DETECTION.ENABLE @@ -255,8 +248,7 @@ def visualize(cfg): slowfast/config/defaults.py """ if cfg.TENSORBOARD.ENABLE and ( - cfg.TENSORBOARD.MODEL_VIS.ENABLE - or cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE + cfg.TENSORBOARD.MODEL_VIS.ENABLE or cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE ): # Set up environment. du.init_distributed_training(cfg) @@ -292,9 +284,7 @@ def visualize(cfg): writer = None if cfg.TENSORBOARD.PREDICTIONS_PATH != "": assert not cfg.DETECTION.ENABLE, "Detection is not supported." - logger.info( - "Visualizing class-level performance from saved results..." - ) + logger.info("Visualizing class-level performance from saved results...") if writer is not None: with pathmgr.open(cfg.TENSORBOARD.PREDICTIONS_PATH, "rb") as f: preds, labels = pickle.load(f, encoding="latin1") @@ -303,20 +293,20 @@ def visualize(cfg): if cfg.TENSORBOARD.MODEL_VIS.ENABLE: if cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.ENABLE: - assert ( - not cfg.DETECTION.ENABLE - ), "Detection task is currently not supported for Grad-CAM visualization." + assert not cfg.DETECTION.ENABLE, ( + "Detection task is currently not supported for Grad-CAM visualization." + ) if cfg.MODEL.ARCH in cfg.MODEL.SINGLE_PATHWAY_ARCH: - assert ( - len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) == 1 - ), "The number of chosen CNN layers must be equal to the number of pathway(s), given {} layer(s).".format( - len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) + assert len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) == 1, ( + "The number of chosen CNN layers must be equal to the number of pathway(s), given {} layer(s).".format( + len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) + ) ) elif cfg.MODEL.ARCH in cfg.MODEL.MULTI_PATHWAY_ARCH: - assert ( - len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) == 2 - ), "The number of chosen CNN layers must be equal to the number of pathway(s), given {} layer(s).".format( - len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) + assert len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) == 2, ( + "The number of chosen CNN layers must be equal to the number of pathway(s), given {} layer(s).".format( + len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) + ) ) else: raise NotImplementedError( @@ -327,17 +317,13 @@ def visualize(cfg): ) ) logger.info( - "Visualize model analysis for {} iterations".format( - len(vis_loader) - ) + "Visualize model analysis for {} iterations".format(len(vis_loader)) ) # Run visualization on the model run_visualization(vis_loader, model, cfg, writer) if cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE: logger.info( - "Visualize Wrong Predictions for {} iterations".format( - len(vis_loader) - ) + "Visualize Wrong Predictions for {} iterations".format(len(vis_loader)) ) perform_wrong_prediction_vis(vis_loader, model, cfg)