[Refactor] Refator ViT (Continue #295)

configs/_base_/datasets/imagenet_bs64_pil_resize_autoaug.py

@@ -0,0 +1,45 @@
+_base_ = [
+    'pipelines/auto_aug.py',
+]
+
+# dataset settings
+dataset_type = 'ImageNet'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='RandomResizedCrop', size=224, backend='pillow'),
+    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),
+    dict(type='AutoAugment', policies={{_base_.policy_imagenet}}),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='ImageToTensor', keys=['img']),
+    dict(type='ToTensor', keys=['gt_label']),
+    dict(type='Collect', keys=['img', 'gt_label'])
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='Resize', size=(256, -1), backend='pillow'),
+    dict(type='CenterCrop', crop_size=224),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='ImageToTensor', keys=['img']),
+    dict(type='Collect', keys=['img'])
+]
+data = dict(
+    samples_per_gpu=64,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        data_prefix='data/imagenet/train',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        data_prefix='data/imagenet/val',
+        ann_file='data/imagenet/meta/val.txt',
+        pipeline=test_pipeline),
+    test=dict(
+        # replace `data/val` with `data/test` for standard test
+        type=dataset_type,
+        data_prefix='data/imagenet/val',
+        ann_file='data/imagenet/meta/val.txt',
+        pipeline=test_pipeline))
+evaluation = dict(interval=1, metric='accuracy')

configs/vision_transformer/vit_base_patch16_224_pretrain_imagenet.py → configs/_base_/datasets/pipelines/auto_aug.py

@@ -1,18 +1,13 @@
-_base_ = [
-    '../_base_/models/vit_base_patch16_224_pretrain.py',
-    '../_base_/datasets/imagenet_bs64_pil_resize.py',
-    '../_base_/schedules/imagenet_bs4096_AdamW.py',
-    '../_base_/default_runtime.py'
-]
-
-policies = [
+# Policy for ImageNet, refers to
+# https://github.com/DeepVoltaire/AutoAugment/blame/master/autoaugment.py
+policy_imagenet = [
     [
         dict(type='Posterize', bits=4, prob=0.4),
         dict(type='Rotate', angle=30., prob=0.6)
     ],
     [
         dict(type='Solarize', thr=256 / 9 * 4, prob=0.6),
-        dict(type='AutoContrast', prob=0.5)
+        dict(type='AutoContrast', prob=0.6)
     ],
     [dict(type='Equalize', prob=0.8),
      dict(type='Equalize', prob=0.6)],
@@ -40,7 +35,7 @@
     ],
     [
         dict(type='Equalize', prob=0.6),
-        dict(type='Posterize', bits=5, prob=0.6)
+        dict(type='Posterize', bits=5, prob=0.4)
     ],
     [
         dict(type='Rotate', angle=30 / 9 * 8, prob=0.8),
@@ -99,45 +94,3 @@
     [dict(type='Equalize', prob=0.8),
      dict(type='Equalize', prob=0.6)],
 ]
-
-dataset_type = 'ImageNet'
-img_norm_cfg = dict(
-    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
-train_pipeline = [
-    dict(type='LoadImageFromFile'),
-    dict(type='RandomResizedCrop', size=224),
-    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),
-    dict(type='AutoAugment', policies=policies),
-    dict(type='Normalize', **img_norm_cfg),
-    dict(type='ImageToTensor', keys=['img']),
-    dict(type='ToTensor', keys=['gt_label']),
-    dict(type='Collect', keys=['img', 'gt_label'])
-]
-
-test_pipeline = [
-    dict(type='LoadImageFromFile'),
-    dict(type='Resize', size=(256, -1), backend='pillow'),
-    dict(type='CenterCrop', crop_size=224),
-    dict(type='Normalize', **img_norm_cfg),
-    dict(type='ImageToTensor', keys=['img']),
-    dict(type='Collect', keys=['img'])
-]
-data = dict(
-    samples_per_gpu=64,
-    workers_per_gpu=2,
-    train=dict(
-        type=dataset_type,
-        data_prefix='data/imagenet/train',
-        pipeline=train_pipeline),
-    val=dict(
-        type=dataset_type,
-        data_prefix='data/imagenet/val',
-        ann_file='data/imagenet/meta/val.txt',
-        pipeline=test_pipeline),
-    test=dict(
-        # replace `data/val` with `data/test` for standard test
-        type=dataset_type,
-        data_prefix='data/imagenet/val',
-        ann_file='data/imagenet/meta/val.txt',
-        pipeline=test_pipeline))
-evaluation = dict(interval=1, metric='accuracy')

configs/_base_/models/vit-base-p16.py

@@ -0,0 +1,25 @@
+# model settings
+model = dict(
+    type='ImageClassifier',
+    backbone=dict(
+        type='VisionTransformer',
+        arch='b',
+        img_size=224,
+        patch_size=16,
+        drop_rate=0.1,
+        init_cfg=[
+            dict(
+                type='Kaiming',
+                layer='Conv2d',
+                mode='fan_in',
+                nonlinearity='linear')
+        ]),
+    neck=None,
+    head=dict(
+        type='VisionTransformerClsHead',
+        num_classes=1000,
+        in_channels=768,
+        loss=dict(
+            type='LabelSmoothLoss', label_smooth_val=0.1,
+            mode='classy_vision'),
+    ))

configs/_base_/models/vit_base_patch32_384_finetune.py → configs/_base_/models/vit-base-p32.py

@@ -3,14 +3,17 @@
     type='ImageClassifier',
     backbone=dict(
         type='VisionTransformer',
-        num_layers=12,
-        embed_dim=768,
-        num_heads=12,
-        img_size=384,
+        arch='b',
+        img_size=224,
         patch_size=32,
-        in_channels=3,
-        feedforward_channels=3072,
-        drop_rate=0.1),
+        drop_rate=0.1,
+        init_cfg=[
+            dict(
+                type='Kaiming',
+                layer='Conv2d',
+                mode='fan_in',
+                nonlinearity='linear')
+        ]),
     neck=None,
     head=dict(
         type='VisionTransformerClsHead',

configs/_base_/models/vit_large_patch16_224_finetune.py → configs/_base_/models/vit-large-p16.py

@@ -3,14 +3,17 @@
     type='ImageClassifier',
     backbone=dict(
         type='VisionTransformer',
-        num_layers=24,
-        embed_dim=1024,
-        num_heads=16,
+        arch='l',
         img_size=224,
         patch_size=16,
-        in_channels=3,
-        feedforward_channels=4096,
-        drop_rate=0.1),
+        drop_rate=0.1,
+        init_cfg=[
+            dict(
+                type='Kaiming',
+                layer='Conv2d',
+                mode='fan_in',
+                nonlinearity='linear')
+        ]),
     neck=None,
     head=dict(
         type='VisionTransformerClsHead',

configs/_base_/models/vit_large_patch32_384_finetune.py → configs/_base_/models/vit-large-p32.py

@@ -3,14 +3,17 @@
     type='ImageClassifier',
     backbone=dict(
         type='VisionTransformer',
-        num_layers=24,
-        embed_dim=1024,
-        num_heads=16,
-        img_size=384,
+        arch='l',
+        img_size=224,
         patch_size=32,
-        in_channels=3,
-        feedforward_channels=4096,
-        drop_rate=0.1),
+        drop_rate=0.1,
+        init_cfg=[
+            dict(
+                type='Kaiming',
+                layer='Conv2d',
+                mode='fan_in',
+                nonlinearity='linear')
+        ]),
     neck=None,
     head=dict(
         type='VisionTransformerClsHead',

configs/vision_transformer/README.md

@@ -14,3 +14,36 @@
   url={https://openreview.net/forum?id=YicbFdNTTy}
 }
 ```
+
+The training step of Vision Transformers is divided into two steps. The first
+step is training the model on a large dataset, like ImageNet-21k, and get the
+pretrain model. And the second step is training the model on the target dataset,
+like ImageNet-1k, and get the finetune model. Here, we provide both pretrain
+models and finetune models.
+
+## Pretrain model
+
+The pre-trained models are converted from [model zoo of Google Research](https://github.com/google-research/vision_transformer#available-vit-models).
+
+### ImageNet 21k
+
+|   Model    | Params(M) |  Flops(G) | Download |
+|:----------:|:---------:|:---------:|:--------:|
+|  ViT-B16\* |   86.86   |   33.03   | [model](https://download.openmmlab.com/mmclassification/v0/vit/pretrain/vit-base-p16_3rdparty_pt-64xb64_in1k-224_20210928-02284250.pth)|
+|  ViT-B32\* |   88.30   |    8.56   | [model](https://download.openmmlab.com/mmclassification/v0/vit/pretrain/vit-base-p32_3rdparty_pt-64xb64_in1k-224_20210928-eee25dd4.pth)|
+|  ViT-L16\* |  304.72   |  116.68   | [model](https://download.openmmlab.com/mmclassification/v0/vit/pretrain/vit-large-p16_3rdparty_pt-64xb64_in1k-224_20210928-0001f9a1.pth)|
+
+*Models with \* are converted from other repos.*
+
+
+## Finetune model
+
+The finetune models are converted from [model zoo of Google Research](https://github.com/google-research/vision_transformer#available-vit-models).
+
+### ImageNet 1k
+|    Model   |  Pretrain    | resolution  | Params(M) |  Flops(G) | Top-1 (%) | Top-5 (%) |   Config   | Download |
+|:----------:|:------------:|:-----------:|:---------:|:---------:|:---------:|:---------:|:----------:|:--------:|
+|  ViT-B16\* | ImageNet-21k |   384x384   |   86.86   |   33.03   |   85.43   |   97.77   | [config](https://github.com/open-mmlab/mmclassification/blob/master/configs/vision_transformer/vit-base-p16_ft-64xb64_in1k-384.py) | [model](https://download.openmmlab.com/mmclassification/v0/vit/finetune/vit-base-p16_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-98e8652b.pth)|
+|  ViT-B32\* | ImageNet-21k |   384x384   |   88.30   |    8.56   |   84.01   |   97.08   | [config](https://github.com/open-mmlab/mmclassification/blob/master/configs/vision_transformer/vit-base-p32_ft-64xb64_in1k-384.py) | [model](https://download.openmmlab.com/mmclassification/v0/vit/finetune/vit-base-p32_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-9cea8599.pth)|
+|  ViT-L16\* | ImageNet-21k |   384x384   |  304.72   |  116.68   |   85.63   |   97.63   | [config](https://github.com/open-mmlab/mmclassification/blob/master/configs/vision_transformer/vit-large-p16_ft-64xb64_in1k-384.py) | [model](https://download.openmmlab.com/mmclassification/v0/vit/finetune/vit-large-p16_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-b20ba619.pth)|
+*Models with \* are converted from other repos.*