Repository: pytorch/vision

Issue:
Update reference scripts to use the "Batteries Included" utils
## 🚀 Feature
As part of the "Batteries Included" initiative (#3911) we are adding a number of new utils that can be used to produce SOTA results. Once those utils are landed, we should update our reference scripts to use them. 

More specifically we need to:
- [x] Update the reference scripts to use PyTorch's new [warmup](https://github.com/pytorch/pytorch/pull/60836) schedulers.
  - The reference scripts should be updated to let users define the `warmup_method` (if any), the `warmup_iters` and the `warmup_factor`. The warmup scheduler should be chained with other existing schedulers.
  - The Object Detection recipe needs to be BC compatible and replace our [custom linear warmup](https://github.com/pytorch/vision/blob/dc42f933f3343c76727dbfba6e4242f1bcb8e1a0/references/detection/engine.py#L20-L24) approach with the one from PyTorch.
  - Similarly the Video Classification recipe needs also to be BC compatible and replace our [customer scheduler](https://github.com/pytorch/vision/blob/master/references/video_classification/scheduler.py#L5).
  - All other recipes (Classification and Segmentation), should be updated to optionally use warm up. The addition should be BC compatible and turned off by default.
- [x] Update the Classification reference to use PyTorch's new label smoothing implementation.
  - We will support the case of label smoothing by using the new `label_smoothing` argument of CrossEntropyLoss (https://github.com/pytorch/pytorch/pull/63122). 
  - A new `label_smoothing` argument should be introduced on the [script](https://github.com/pytorch/vision/blob/c2ab0c59f42babf9ad01aa616cd8a901daac86dd/references/classification/train.py#L165). To ensure BC, its default value should be 0.0 (no label smoothing).
- [x] Update the Classification reference to use Mixup and Cutmix.
  - After implementing the Mixup and Cutmix augmentations, update the reference script of classification to use them on training. https://github.com/pytorch/vision/pull/4379 
  - The new `mixup-alpha` and `cutmix-alpha` arguments should have 0.0 default values to ensure BC.
- [x] Update the Classification reference to use [EMA](https://github.com/pytorch/vision/issues/4346).
  - Add Exponential Moving Average support on the Classification scripts by extending the `torch.optim.swa_utils.AveragedModel` util and ensure the EMA model is evaluated at the end of each epoch.
  - Use a feature switch `model-ema` param to ensure BC and ensure we have a good `model-ema-decay` default value.

Improve the accuracy of Classification models by using SOTA recipes and primitives
## 🚀 Feature
Update the weights of all pre-trained models to improve their accuracy.

## Motivation

<h4 id="new-recipe-with-fixres">New Recipe + FixRes mitigations</h4>

```
torchrun --nproc_per_node=8 train.py --model $MODEL_NAME --batch-size 128 --lr 0.5 \
--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.00002 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--train-crop-size 176 --model-ema --val-resize-size 232
```

Using a recipe which includes Warmup, Cosine Annealing, Label Smoothing, Mixup, Cutmix, Random Erasing, TrivialAugment, No BN weight decay, EMA and long training cycles and optional FixRes mitigations we are able to improve the `resnet50` accuracy by over 4.5 points. For more information on the training recipe, check [here](https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/):
```
Old ResNet50:
Acc@1 76.130 Acc@5 92.862

New ResNet50:
Acc@1 80.674 Acc@5 95.166
```

Running other models through the same recipe, achieves the following improved accuracies:
```
ResNet101:
Acc@1 81.728 Acc@5 95.670

ResNet152:
Acc@1 82.042 Acc@5 95.926

ResNeXt50_32x4d:
Acc@1 81.116 Acc@5 95.478

ResNeXt101_32x8d:
Acc@1 82.834 Acc@5 96.228

MobileNetV3 Large:
Acc@1 74.938 Acc@5 92.496

Wide ResNet50 2:
Acc@1 81.602 Acc@5 95.758 (@prabhat00155)

Wide ResNet101 2:
Acc@1 82.492 Acc@5 96.110 (@prabhat00155)

regnet_x_400mf:
Acc@1 74.864 Acc@5 92.322 (@kazhang)

regnet_x_800mf:
Acc@1 77.522 Acc@5 93.826 (@kazhang)

regnet_x_1_6gf:
Acc@1 79.668 Acc@5 94.922 (@kazhang)
```

<h4 id="new-recipe">New Recipe (without FixRes mitigations)</h4>

```
torchrun --nproc_per_node=8 train.py --model $MODEL_NAME --batch-size 128 --lr 0.5 \
--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.00002 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--model-ema --val-resize-size 232
```

Removing the optional FixRes mitigations seems to yield better results for some deeper architectures and variants with larger receptive fields:
```
ResNet101:
Acc@1 81.886 Acc@5 95.780

ResNet152:
Acc@1 82.284 Acc@5 96.002

ResNeXt50_32x4d:
Acc@1 81.198 Acc@5 95.340

ResNeXt101_32x8d:
Acc@1 82.812 Acc@5 96.226

MobileNetV3 Large:
Acc@1 75.152 Acc@5 92.634

Wide ResNet50_2:
Acc@1 81.452 Acc@5 95.544 (@prabhat00155)

Wide ResNet101_2:
Acc@1 82.510 Acc@5 96.020 (@prabhat00155)

regnet_x_3_2gf:
Acc@1 81.196 Acc@5 95.430

regnet_x_8gf:
Acc@1 81.682 Acc@5 95.678

regnet_x_16g:
Acc@1 82.716 Acc@5 96.196

regnet_x_32gf:
Acc@1 83.014 Acc@5 96.288

regnet_y_400mf:
Acc@1 75.804 Acc@5 92.742

regnet_y_800mf:
Acc@1 78.828 Acc@5 94.502

regnet_y_1_6gf:
Acc@1 80.876 Acc@5 95.444

regnet_y_3_2gf:
Acc@1 81.982 Acc@5 95.972

regnet_y_8gf:
Acc@1 82.828 Acc@5 96.330

regnet_y_16gf:
Acc@1 82.886 Acc@5 96.328

regnet_y_32gf:
Acc@1 83.368 Acc@5 96.498
```

<h4 id="new-recipe-with-reg-tuning">New Recipe + Regularization tuning</h4>

```
torchrun --nproc_per_node=8 train.py --model $MODEL_NAME --batch-size 128 --lr 0.5 \
--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.00001 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--model-ema --val-resize-size 232
```

Adjusting slightly the regularization can help us improve the following:
```
MobileNetV3 Large:
Acc@1 75.274 Acc@5 92.566
```
In addition to regularization adjustment we can also apply the Repeated Augmentation trick ` --ra-sampler --ra-reps 4`:

```
MobileNetV2:
Acc@1 72.154 Acc@5 90.822
```

<h4 id="ptq-models">Post-Training Quantized models</h4>

```
ResNet50:
Acc@1 80.282 Acc@5 94.976

ResNeXt101_32x8d:
Acc@1 82.574 Acc@5 96.132
```

<h4 id="new-recipe-with-lr-wd-crop-tuning">New Recipe (LR+weight_decay+train_crop_size tuning)</h4>

```
torchrun --ngpus 8 --nodes 1 --model $MODEL_NAME --batch-size 128 --lr 1 \
--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.000002 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--train-crop-size 208 --model-ema --val-crop-size 240 --val-resize-size 255
```

```
EfficientNet-B1:
Acc@1 79.838 Acc@5 94.934
```

## Pitch

To be able to improve the pre-trained model accuracy, we need to complete the "Batteries Included" work as #3911. Moreover we will need to extend our existing model builders to support multiple weights as described at #4611. Then we will be able to:
- Update our reference scripts for classification to support the new primitives added by the "Batteries Included" initiative.
- Find a good training recipe for the most important pre-trained models and re-train them. Note that different training configuration might be required for different types of models (for example mobile models are less likely to overfit comparing to bigger models and thus make use of different recipes/primitives)
- Update the weights of the models in the library.

cc @datumbox @vfdev-5


Discussion:

Resolving patch:
diff --git a/references/classification/train.py b/references/classification/train.py
--- a/references/classification/train.py
+++ b/references/classification/train.py
@@ -186,17 +186,19 @@ def main(args):
         sampler=test_sampler, num_workers=args.workers, pin_memory=True)
 
     print("Creating model")
-    model = torchvision.models.__dict__[args.model](pretrained=args.pretrained)
+    model = torchvision.models.__dict__[args.model](pretrained=args.pretrained, num_classes=num_classes)
     model.to(device)
+
     if args.distributed and args.sync_bn:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
 
     criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing)
 
     opt_name = args.opt.lower()
-    if opt_name == 'sgd':
+    if opt_name.startswith("sgd"):
         optimizer = torch.optim.SGD(
-            model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+            model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay,
+            nesterov="nesterov" in opt_name)
     elif opt_name == 'rmsprop':
         optimizer = torch.optim.RMSprop(model.parameters(), lr=args.lr, momentum=args.momentum,
                                         weight_decay=args.weight_decay, eps=0.0316, alpha=0.9)
@@ -214,15 +216,25 @@ def main(args):
     elif args.lr_scheduler == 'cosineannealinglr':
         main_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
                                                                        T_max=args.epochs - args.lr_warmup_epochs)
+    elif args.lr_scheduler == 'exponentiallr':
+        main_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=args.lr_gamma)
     else:
-        raise RuntimeError("Invalid lr scheduler '{}'. Only StepLR and CosineAnnealingLR "
+        raise RuntimeError("Invalid lr scheduler '{}'. Only StepLR, CosineAnnealingLR and ExponentialLR "
                            "are supported.".format(args.lr_scheduler))
 
     if args.lr_warmup_epochs > 0:
+        if args.lr_warmup_method == 'linear':
+            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=args.lr_warmup_decay,
+                                                                    total_iters=args.lr_warmup_epochs)
+        elif args.lr_warmup_method == 'constant':
+            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=args.lr_warmup_decay,
+                                                                      total_iters=args.lr_warmup_epochs)
+        else:
+            raise RuntimeError(f"Invalid warmup lr method '{args.lr_warmup_method}'. Only linear and constant "
+                               "are supported.")
         lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
             optimizer,
-            schedulers=[torch.optim.lr_scheduler.ConstantLR(optimizer, factor=args.lr_warmup_decay,
-                                                            total_iters=args.lr_warmup_epochs), main_lr_scheduler],
+            schedulers=[warmup_lr_scheduler, main_lr_scheduler],
             milestones=[args.lr_warmup_epochs]
         )
     else:
@@ -307,7 +319,9 @@ def get_args_parser(add_help=True):
     parser.add_argument('--cutmix-alpha', default=0.0, type=float, help='cutmix alpha (default: 0.0)')
     parser.add_argument('--lr-scheduler', default="steplr", help='the lr scheduler (default: steplr)')
     parser.add_argument('--lr-warmup-epochs', default=0, type=int, help='the number of epochs to warmup (default: 0)')
-    parser.add_argument('--lr-warmup-decay', default=0.01, type=int, help='the decay for lr')
+    parser.add_argument('--lr-warmup-method', default="constant", type=str,
+                        help='the warmup method (default: constant)')
+    parser.add_argument('--lr-warmup-decay', default=0.01, type=float, help='the decay for lr')
     parser.add_argument('--lr-step-size', default=30, type=int, help='decrease lr every step-size epochs')
     parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')
     parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
diff --git a/references/segmentation/train.py b/references/segmentation/train.py
--- a/references/segmentation/train.py
+++ b/references/segmentation/train.py
@@ -220,7 +220,7 @@ def get_args_parser(add_help=True):
                         dest='weight_decay')
     parser.add_argument('--lr-warmup-epochs', default=0, type=int, help='the number of epochs to warmup (default: 0)')
     parser.add_argument('--lr-warmup-method', default="linear", type=str, help='the warmup method (default: linear)')
-    parser.add_argument('--lr-warmup-decay', default=0.01, type=int, help='the decay for lr')
+    parser.add_argument('--lr-warmup-decay', default=0.01, type=float, help='the decay for lr')
     parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
     parser.add_argument('--output-dir', default='.', help='path where to save')
     parser.add_argument('--resume', default='', help='resume from checkpoint')
diff --git a/references/video_classification/train.py b/references/video_classification/train.py
--- a/references/video_classification/train.py
+++ b/references/video_classification/train.py
@@ -296,7 +296,7 @@ def parse_args():
     parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')
     parser.add_argument('--lr-warmup-epochs', default=10, type=int, help='the number of epochs to warmup (default: 10)')
     parser.add_argument('--lr-warmup-method', default="linear", type=str, help='the warmup method (default: linear)')
-    parser.add_argument('--lr-warmup-decay', default=0.001, type=int, help='the decay for lr')
+    parser.add_argument('--lr-warmup-decay', default=0.001, type=float, help='the decay for lr')
     parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
     parser.add_argument('--output-dir', default='.', help='path where to save')
     parser.add_argument('--resume', default='', help='resume from checkpoint')


Test patch: