Merge branch 'main' into transforms/mixupcutmix

Author: datumbox

docs/source/feature_extraction.rst

@@ -39,6 +39,7 @@ Here is an example of how we might extract features for MaskRCNN:
   from torchvision.models.feature_extraction import get_graph_node_names
   from torchvision.models.feature_extraction import create_feature_extractor
   from torchvision.models.detection.mask_rcnn import MaskRCNN
+  from torchvision.models.detection.backbone_utils import LastLevelMaxPool
   from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork
 
 
@@ -57,7 +58,7 @@ Here is an example of how we might extract features for MaskRCNN:
   # that appears in each of the main layers:
   return_nodes = {
       # node_name: user-specified key for output dict
-      'layer1.2.relu_2': 'layer1',  
+      'layer1.2.relu_2': 'layer1',
       'layer2.3.relu_2': 'layer2',
       'layer3.5.relu_2': 'layer3',
       'layer4.2.relu_2': 'layer4',
@@ -70,7 +71,7 @@ Here is an example of how we might extract features for MaskRCNN:
   # performed is the one that corresponds to the output you desire. You should
   # consult the source code for the input model to confirm.)
   return_nodes = {
-      'layer1': 'layer1',  
+      'layer1': 'layer1',
       'layer2': 'layer2',
       'layer3': 'layer3',
       'layer4': 'layer4',
@@ -79,7 +80,7 @@ Here is an example of how we might extract features for MaskRCNN:
   # Now you can build the feature extractor. This returns a module whose forward
   # method returns a dictionary like:
   # {
-  #     'layer1': ouput of layer 1,  
+  #     'layer1': ouput of layer 1,
   #     'layer2': ouput of layer 2,
   #     'layer3': ouput of layer 3,
   #     'layer4': ouput of layer 4,
@@ -94,10 +95,11 @@ Here is an example of how we might extract features for MaskRCNN:
           super(Resnet50WithFPN, self).__init__()
           # Get a resnet50 backbone
           m = resnet50()
-          # Extract 4 main layers (note: you can also provide a list for return
-          # nodes if the keys and the values are the same)
+          # Extract 4 main layers (note: MaskRCNN needs this particular name
+          # mapping for return nodes)
           self.body = create_feature_extractor(
-              m, return_nodes=['layer1', 'layer2', 'layer3', 'layer4'])
+              m, return_nodes={f'layer{k}': str(v)
+                               for v, k in enumerate([1, 2, 3, 4])})
           # Dry run to get number of channels for FPN
           inp = torch.randn(2, 3, 224, 224)
           with torch.no_grad():
@@ -106,7 +108,8 @@ Here is an example of how we might extract features for MaskRCNN:
           # Build FPN
           self.out_channels = 256
           self.fpn = FeaturePyramidNetwork(
-              in_channels_list, out_channels=self.out_channels)
+              in_channels_list, out_channels=self.out_channels,
+              extra_blocks=LastLevelMaxPool())
 
       def forward(self, x):
           x = self.body(x)

references/classification/train.py

@@ -17,7 +17,8 @@
     amp = None
 
 
-def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, print_freq, apex=False):
+def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch,
+                    print_freq, apex=False, model_ema=None):
     model.train()
     metric_logger = utils.MetricLogger(delimiter="  ")
     metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))
@@ -45,11 +46,14 @@ def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, pri
         metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
         metric_logger.meters['img/s'].update(batch_size / (time.time() - start_time))
 
+    if model_ema:
+        model_ema.update_parameters(model)
 
-def evaluate(model, criterion, data_loader, device, print_freq=100):
+
+def evaluate(model, criterion, data_loader, device, print_freq=100, log_suffix=''):
     model.eval()
     metric_logger = utils.MetricLogger(delimiter="  ")
-    header = 'Test:'
+    header = f'Test: {log_suffix}'
     with torch.no_grad():
         for image, target in metric_logger.log_every(data_loader, print_freq, header):
             image = image.to(device, non_blocking=True)
@@ -199,12 +203,18 @@ def main(args):
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
         model_without_ddp = model.module
 
+    model_ema = None
+    if args.model_ema:
+        model_ema = utils.ExponentialMovingAverage(model_without_ddp, device=device, decay=args.model_ema_decay)
+
     if args.resume:
         checkpoint = torch.load(args.resume, map_location='cpu')
         model_without_ddp.load_state_dict(checkpoint['model'])
         optimizer.load_state_dict(checkpoint['optimizer'])
         lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
         args.start_epoch = checkpoint['epoch'] + 1
+        if model_ema:
+            model_ema.load_state_dict(checkpoint['model_ema'])
 
     if args.test_only:
         evaluate(model, criterion, data_loader_test, device=device)
@@ -215,16 +225,20 @@ def main(args):
     for epoch in range(args.start_epoch, args.epochs):
         if args.distributed:
             train_sampler.set_epoch(epoch)
-        train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args.print_freq, args.apex)
+        train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args.print_freq, args.apex, model_ema)
         lr_scheduler.step()
         evaluate(model, criterion, data_loader_test, device=device)
+        if model_ema:
+            evaluate(model_ema, criterion, data_loader_test, device=device, log_suffix='EMA')
         if args.output_dir:
             checkpoint = {
                 'model': model_without_ddp.state_dict(),
                 'optimizer': optimizer.state_dict(),
                 'lr_scheduler': lr_scheduler.state_dict(),
                 'epoch': epoch,
                 'args': args}
+            if model_ema:
+                checkpoint['model_ema'] = model_ema.state_dict()
             utils.save_on_master(
                 checkpoint,
                 os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
@@ -306,6 +320,12 @@ def get_args_parser(add_help=True):
     parser.add_argument('--world-size', default=1, type=int,
                         help='number of distributed processes')
     parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
+    parser.add_argument(
+        '--model-ema', action='store_true',
+        help='enable tracking Exponential Moving Average of model parameters')
+    parser.add_argument(
+        '--model-ema-decay', type=float, default=0.99,
+        help='decay factor for Exponential Moving Average of model parameters(default: 0.99)')
 
     return parser
 

references/classification/utils.py

@@ -161,6 +161,18 @@ def log_every(self, iterable, print_freq, header=None):
         print('{} Total time: {}'.format(header, total_time_str))
 
 
+class ExponentialMovingAverage(torch.optim.swa_utils.AveragedModel):
+    """Maintains moving averages of model parameters using an exponential decay.
+    ``ema_avg = decay * avg_model_param + (1 - decay) * model_param``
+    `torch.optim.swa_utils.AveragedModel <https://pytorch.org/docs/stable/optim.html#custom-averaging-strategies>`_
+    is used to compute the EMA.
+    """
+    def __init__(self, model, decay, device='cpu'):
+        ema_avg = (lambda avg_model_param, model_param, num_averaged:
+                   decay * avg_model_param + (1 - decay) * model_param)
+        super().__init__(model, device, ema_avg)
+
+
 def accuracy(output, target, topk=(1,)):
     """Computes the accuracy over the k top predictions for the specified values of k"""
     with torch.no_grad():

test/cpp/test_custom_operators.cpp

@@ -18,7 +18,7 @@ TEST(test_custom_operators, nms) {
   double thresh = 0.7;
 
   torch::jit::push(stack, boxes, scores, thresh);
-  op->getOperation()(&stack);
+  op->getOperation()(stack);
   at::Tensor output_jit;
   torch::jit::pop(stack, output_jit);
 
@@ -47,7 +47,7 @@ TEST(test_custom_operators, roi_align_visible) {
   bool aligned = true;
 
   torch::jit::push(stack, input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned);
-  op->getOperation()(&stack);
+  op->getOperation()(stack);
   at::Tensor output_jit;
   torch::jit::pop(stack, output_jit);
 

test/test_backbone_utils.py

@@ -39,7 +39,7 @@ class TestFxFeatureExtraction:
         'num_classes': 1,
         'pretrained': False
     }
-    leaf_modules = [torchvision.ops.StochasticDepth]
+    leaf_modules = []
 
     def _create_feature_extractor(self, *args, **kwargs):
         """

test/test_datasets.py

@@ -512,7 +512,7 @@ def inject_fake_data(self, tmpdir, config):
         return dict(num_examples=num_images_per_split[config["split"]], attr_names=attr_names)
 
     def _create_split_txt(self, root):
-        num_images_per_split = dict(train=3, valid=2, test=1)
+        num_images_per_split = dict(train=4, valid=3, test=2)
 
         data = [
             [self._SPLIT_TO_IDX[split]] for split, num_images in num_images_per_split.items() for _ in range(num_images)
@@ -595,6 +595,17 @@ def test_attr_names(self):
         with self.create_dataset() as (dataset, info):
             assert tuple(dataset.attr_names) == info["attr_names"]
 
+    def test_images_names_split(self):
+        with self.create_dataset(split='all') as (dataset, _):
+            all_imgs_names = set(dataset.filename)
+
+        merged_imgs_names = set()
+        for split in ["train", "valid", "test"]:
+            with self.create_dataset(split=split) as (dataset, _):
+                merged_imgs_names.update(dataset.filename)
+
+        assert merged_imgs_names == all_imgs_names
+
 
 class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.VOCSegmentation

torchvision/datasets/celeba.py

@@ -99,7 +99,10 @@ def __init__(
 
         mask = slice(None) if split_ is None else (splits.data == split_).squeeze()
 
-        self.filename = splits.index
+        if mask == slice(None):  # if split == "all"
+            self.filename = splits.index
+        else:
+            self.filename = [splits.index[i] for i in torch.squeeze(torch.nonzero(mask))]
         self.identity = identity.data[mask]
         self.bbox = bbox.data[mask]
         self.landmarks_align = landmarks_align.data[mask]

torchvision/models/detection/faster_rcnn.py

@@ -300,6 +300,9 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
     """
     Constructs a Faster R-CNN model with a ResNet-50-FPN backbone.
 
+    Reference: `"Faster R-CNN: Towards Real-Time Object Detection with
+    Region Proposal Networks" <https://arxiv.org/abs/1506.01497>`_.
+
     The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
     image, and should be in ``0-1`` range. Different images can have different sizes.
 

torchvision/models/detection/keypoint_rcnn.py

@@ -278,6 +278,8 @@ def keypointrcnn_resnet50_fpn(pretrained=False, progress=True,
     """
     Constructs a Keypoint R-CNN model with a ResNet-50-FPN backbone.
 
+    Reference: `"Mask R-CNN" <https://arxiv.org/abs/1703.06870>`_.
+
     The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
     image, and should be in ``0-1`` range. Different images can have different sizes.
 

torchvision/models/detection/mask_rcnn.py

@@ -271,6 +271,8 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True,
     """
     Constructs a Mask R-CNN model with a ResNet-50-FPN backbone.
 
+    Reference: `"Mask R-CNN" <https://arxiv.org/abs/1703.06870>`_.
+
     The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
     image, and should be in ``0-1`` range. Different images can have different sizes.