Adding Swin Transformer architecture (#5491)

xiaohu2015 · datumbox · jdsgomes · web-flow · commit e288f6ca01a5 · 2022-04-27T12:03:21.000+01:00
* add swin transformer * Update swin_transformer.py * Update swin_transformer.py * fix lint * fix lint * refactor code * add swin_transformer * Update swin_transformer.py * fix bug * refactor code * fix lint * update init_weights * move shift_window into attention * refactor code * fix bug * Update swin_transformer.py * Update swin_transformer.py * fix lint * add patch_merge * fix bug * Update swin_transformer.py * Update swin_transformer.py * Update swin_transformer.py * refactor code * Update swin_transformer.py * refactor code * fix lint * refactor code * add swin_tiny * add swin_tiny.pkl * fix lint * Delete ModelTester.test_swin_tiny_expect.pkl * add swin_tiny * add * add Optional to bias * update init weights * update init_weights and add no weight decay * add no weight decay * add set_weight_decay * add set_weight_decay * fix lint * fix lint * add lr_cos_min * add other swin models * Update torchvision/models/swin_transformer.py Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com> * refactor doc * Update utils.py * Update train.py * Update train.py * Update swin_transformer.py * update model builder * fix lint * add * Update torchvision/models/swin_transformer.py Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com> * Update torchvision/models/swin_transformer.py Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com> * update other model * simplify the model name just like ViT * add lr_cos_min * fix lint * fix lint * Update swin_transformer.py * Update swin_transformer.py * Update swin_transformer.py * Delete ModelTester.test_swin_tiny_expect.pkl * add swin_t * refactor code * Update train.py * add swin_s * ignore a error of mypy * Update swin_transformer.py * fix lint * add swin_b * add swin_l * refactor code * Update train.py * move relative_position_bias to __init__ * fix formatting * Revert "fix formatting" This reverts commit 41faba2. * Revert "move relative_position_bias to __init__" This reverts commit f061544. * refactor code * Remove deprecated meta-data from `_COMMON_META` * fix linter * add pretrained weights for swin_t * fix format * apply ufmt * add documentation * update references README * adding new style docs * update pre-trained weights values * remove other variants * fix typo * Remove expect for the variants not yet supported Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com> Co-authored-by: Joao Gomes <jdsgomes@fb.com>
diff --git a/docs/source/models.rst b/docs/source/models.rst
@@ -42,6 +42,7 @@ architectures for image classification:
 -  `RegNet`_
 -  `VisionTransformer`_
 -  `ConvNeXt`_
+-  `SwinTransformer`_
 
 You can construct a model with random weights by calling its constructor:
 
@@ -97,6 +98,7 @@ You can construct a model with random weights by calling its constructor:
     convnext_small = models.convnext_small()
     convnext_base = models.convnext_base()
     convnext_large = models.convnext_large()
+    swin_t = models.swin_t()
 
 We provide pre-trained models, using the PyTorch :mod:`torch.utils.model_zoo`.
 
@@ -219,6 +221,7 @@ convnext_tiny                     82.520          96.146
 convnext_small                    83.616          96.650
 convnext_base                     84.062          96.870
 convnext_large                    84.414          96.976
+swin_t                            81.358          95.526
 ================================  =============   =============
 
 
@@ -238,6 +241,7 @@ convnext_large                    84.414          96.976
 .. _RegNet: https://arxiv.org/abs/2003.13678
 .. _VisionTransformer: https://arxiv.org/abs/2010.11929
 .. _ConvNeXt: https://arxiv.org/abs/2201.03545
+.. _SwinTransformer: https://arxiv.org/abs/2103.14030
 
 .. currentmodule:: torchvision.models
 
@@ -450,6 +454,15 @@ ConvNeXt
     convnext_base
     convnext_large
 
+SwinTransformer
+--------
+
+.. autosummary::
+    :toctree: generated/
+    :template: function.rst
+
+    swin_t
+
 Quantized Models
 ----------------
 
diff --git a/docs/source/models/swin_transformer.rst b/docs/source/models/swin_transformer.rst
@@ -0,0 +1,25 @@
+SwinTransformer
+===============
+
+.. currentmodule:: torchvision.models
+
+The SwinTransformer model is based on the `Swin Transformer: Hierarchical Vision 
+Transformer using Shifted Windows <https://arxiv.org/abs/2103.14030>`__
+paper.
+
+
+Model builders
+--------------
+
+The following model builders can be used to instanciate an SwinTransformer model. 
+`swin_t` can be instantiated with pre-trained weights and all others without. 
+All the model builders internally rely on the ``torchvision.models.swin_transformer.SwinTransformer`` 
+base class. Please refer to the `source code
+<https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py>`_ for
+more details about this class.
+
+.. autosummary::
+    :toctree: generated/
+    :template: function.rst
+
+    swin_t
diff --git a/docs/source/models_new.rst b/docs/source/models_new.rst
@@ -46,6 +46,7 @@ weights:
    models/resnet
    models/resnext
    models/squeezenet
+   models/swin_transformer
    models/vgg
    models/vision_transformer
 
diff --git a/references/classification/README.md b/references/classification/README.md
@@ -224,6 +224,18 @@ Note that the above command corresponds to training on a single node with 8 GPUs
 For generatring the pre-trained weights, we trained with 2 nodes, each with 8 GPUs (for a total of 16 GPUs),
 and `--batch_size 64`.
 
+
+### SwinTransformer
+```
+torchrun --nproc_per_node=8 train.py\ 
+--model swin_t  --epochs 300 --batch-size 128 --opt adamw --lr 0.001 --weight-decay 0.05 --norm-weight-decay 0.0\
+--bias-weight-decay 0.0 --transformer-embedding-decay 0.0 --lr-scheduler cosineannealinglr --lr-min 0.00001 --lr-warmup-method linear\
+--lr-warmup-epochs 20 --lr-warmup-decay 0.01 --amp --label-smoothing 0.1 --mixup-alpha 0.8\
+--clip-grad-norm 5.0 --cutmix-alpha 1.0 --random-erase 0.25 --interpolation bicubic --auto-augment ra
+```
+Note that `--val-resize-size` was optimized in a post-training step, see their `Weights` entry for the exact value.
+
+
 ## Mixed precision training
 Automatic Mixed Precision (AMP) training on GPU for Pytorch can be enabled with the [torch.cuda.amp](https://pytorch.org/docs/stable/amp.html?highlight=amp#module-torch.cuda.amp).
 
diff --git a/references/classification/train.py b/references/classification/train.py
@@ -233,7 +233,7 @@ def main(args):
     if args.bias_weight_decay is not None:
         custom_keys_weight_decay.append(("bias", args.bias_weight_decay))
     if args.transformer_embedding_decay is not None:
-        for key in ["class_token", "position_embedding", "relative_position_bias"]:
+        for key in ["class_token", "position_embedding", "relative_position_bias_table"]:
             custom_keys_weight_decay.append((key, args.transformer_embedding_decay))
     parameters = utils.set_weight_decay(
         model,
@@ -267,7 +267,7 @@ def main(args):
         main_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
     elif args.lr_scheduler == "cosineannealinglr":
         main_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
-            optimizer, T_max=args.epochs - args.lr_warmup_epochs
+            optimizer, T_max=args.epochs - args.lr_warmup_epochs, eta_min=args.lr_min
         )
     elif args.lr_scheduler == "exponentiallr":
         main_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=args.lr_gamma)
@@ -424,6 +424,7 @@ def get_args_parser(add_help=True):
     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("--lr-min", default=0.0, type=float, help="minimum lr of lr schedule (default: 0.0)")
     parser.add_argument("--print-freq", default=10, type=int, help="print frequency")
     parser.add_argument("--output-dir", default=".", type=str, help="path to save outputs")
     parser.add_argument("--resume", default="", type=str, help="path of checkpoint")
diff --git a/test/expect/ModelTester.test_swin_t_expect.pkl b/test/expect/ModelTester.test_swin_t_expect.pkl
diff --git a/torchvision/models/__init__.py b/torchvision/models/__init__.py
@@ -12,6 +12,7 @@
 from .squeezenet import *
 from .vgg import *
 from .vision_transformer import *
+from .swin_transformer import *
 from . import detection
 from . import optical_flow
 from . import quantization
diff --git a/torchvision/models/swin_transformer.py b/torchvision/models/swin_transformer.py
