Rebuild vision transformer model

docs/source/changelog.md

@@ -4,22 +4,25 @@
 
 **New Features**
 
-- Support `trunc_normal_` in `flowvision.layers.weight_init` (#92)
-- 
+- Support `trunc_normal_` in `flowvision.layers.weight_init` [#92](https://github.com/Oneflow-Inc/vision/pull/92)
+- Support [DeiT](https://arxiv.org/abs/2012.12877) model [#115](https://github.com/Oneflow-Inc/vision/pull/115)
 
 **Bug Fixes**
 
-- Fix `EfficientNet` and `Res2Net`(#122)
-- 
+- Fix `F.normalize` usage in SSD [#116](https://github.com/Oneflow-Inc/vision/pull/116)
+- Fix bug in `EfficientNet` and `Res2Net` [#122](https://github.com/Oneflow-Inc/vision/pull/122)
 
 **Improvements**
 
-- Refine `trunc_normal_` and `linspace` usage in `Swin-T`, `CrossFormer`, `PVT` and `CSwin` models. (#100)
--
+- Refator `trunc_normal_` and `linspace` usage in Swin-T, Cross-Former, PVT and CSWin models. [#100](https://github.com/Oneflow-Inc/vision/pull/100)
+- Refator `Vision Transformer` model [#115](https://github.com/Oneflow-Inc/vision/pull/115)
 
-**Contributors**
 
-A total of xx developers contributed to this release.
+**Docs Update**
+- Update `Vision Transformer` docs [#115](https://github.com/Oneflow-Inc/vision/pull/115)
+
+
+**Contributors**
 
-Thanks @rentainhe, @simonJJJ, @kaijieshi7, @lixiang007666, @Ldpe2G
+A total of 5 developers contributed to this release. Thanks @rentainhe, @simonJJJ, @kaijieshi7, @lixiang007666, @Ldpe2G
 

docs/source/flowvision.models.rst

@@ -26,6 +26,7 @@ architectures for image classification:
 -  `EfficientNet`_
 -  `ReXNet`_
 -  `ViT`_
+-  `DeiT`_
 -  `PVT`_
 -  `Swin-Transformer`_
 -  `CSwin-Transformer`_
@@ -52,6 +53,7 @@ architectures for image classification:
 .. _MNASNet: https://arxiv.org/abs/1807.11626
 .. _GhostNet: https://arxiv.org/abs/1911.11907
 .. _ViT: https://arxiv.org/abs/2010.11929
+.. _DeiT: https://arxiv.org/abs/2012.12877
 .. _PVT: https://arxiv.org/abs/2102.12122
 .. _ResMLP: https://arxiv.org/abs/2105.03404
 .. _Swin-Transformer: https://arxiv.org/abs/2103.14030
@@ -218,12 +220,52 @@ ViT
 ------
 .. automodule:: flowvision.models
     :members: 
-        vit_b_16_224,
-        vit_b_16_384,
-        vit_b_32_224,
-        vit_b_32_384,
-        vit_l_16_384,
-        vit_l_32_384,
+        vit_tiny_patch16_224,
+        vit_tiny_patch16_384,
+        vit_small_patch32_224,
+        vit_small_patch32_384,
+        vit_small_patch16_224,
+        vit_small_patch16_384,
+        vit_base_patch32_224,
+        vit_base_patch32_384,
+        vit_base_patch16_224,
+        vit_base_patch16_384,
+        vit_base_patch8_224,
+        vit_large_patch32_224,
+        vit_large_patch32_384,
+        vit_large_patch16_224,
+        vit_large_patch16_384,
+        vit_base_patch16_sam_224,
+        vit_base_patch32_sam_224,
+        vit_huge_patch14_224,
+        vit_giant_patch14_224,
+        vit_gigantic_patch14_224,
+        vit_tiny_patch16_224_in21k,
+        vit_small_patch32_224_in21k,
+        vit_small_patch16_224_in21k,
+        vit_base_patch32_224_in21k,
+        vit_base_patch16_224_in21k,
+        vit_base_patch8_224_in21k,
+        vit_large_patch32_224_in21k,
+        vit_large_patch16_224_in21k,
+        vit_huge_patch14_224_in21k,
+        vit_base_patch16_224_miil_in21k,
+        vit_base_patch16_224_miil,
+
+
+DeiT
+------
+.. automodule:: flowvision.models
+    :members: 
+        deit_tiny_patch16_224,
+        deit_small_patch16_224,
+        deit_base_patch16_224,
+        deit_base_patch16_384,
+        deit_tiny_distilled_patch16_224,
+        deit_small_distilled_patch16_224,
+        deit_base_distilled_patch16_224,
+        deit_base_distilled_patch16_384,
+
 
 PVT
 ------

flowvision/models/__init__.py

@@ -1,6 +1,5 @@
 from .alexnet import *
 from .densenet import *
-from .vit import *
 from .vgg import *
 from .mnasnet import *
 from .resnet import *
@@ -22,6 +21,7 @@
 from .ghostnet import *
 from .res2net import *
 from .efficientnet import *
+from .vision_transformer import *
 
 from . import neural_style_transfer
 from . import detection

flowvision/models/helpers.py

@@ -1,8 +1,11 @@
-""" Layer/Module Helpers
-Hacked together by / Copyright 2020 Ross Wightman
+"""
+Modified from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/helpers.py
 """
 from itertools import repeat
 import collections.abc
+from typing import Any, Callable, Optional, Tuple
+
+import oneflow.nn as nn
 
 
 # From PyTorch internals
@@ -29,3 +32,37 @@ def make_divisible(v, divisor=8, min_value=None, round_limit=0.9):
     if new_v < round_limit * v:
         new_v += divisor
     return new_v
+
+
+def named_apply(
+    fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False
+) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(
+            fn=fn,
+            module=child_module,
+            name=child_name,
+            depth_first=depth_first,
+            include_root=True,
+        )
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+def named_modules(module: nn.Module, name="", depth_first=True, include_root=False):
+    if not depth_first and include_root:
+        yield name, module
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        yield from named_modules(
+            module=child_module,
+            name=child_name,
+            depth_first=depth_first,
+            include_root=True,
+        )
+    if depth_first and include_root:
+        yield name, module

flowvision/models/mlp_mixer.py

@@ -5,12 +5,12 @@
 import oneflow.nn as nn
 import oneflow.nn.init as init
 
-from .utils import load_state_dict_from_url, named_apply
-from .registry import ModelCreator
 from flowvision.layers.regularization import DropPath
 from flowvision.layers.blocks import PatchEmbed
 from flowvision.layers.weight_init import lecun_normal_
-
+from .utils import load_state_dict_from_url
+from .helpers import named_apply
+from .registry import ModelCreator
 
 model_urls = {
     "mlp_mixer_s16_224": None,

flowvision/models/utils.py

@@ -11,10 +11,8 @@
 from urllib.parse import urlparse
 from urllib.request import Request, urlopen
 from tqdm import tqdm
-from typing import Any, Callable, Optional, Tuple
 
 import oneflow as flow
-import oneflow.nn as nn
 
 HASH_REGEX = re.compile(r"([a-f0-9]*)_")
 
@@ -184,37 +182,3 @@ def download_url_to_file(url, dst, hash_prefix=None, progress=True):
         f.close()
         if os.path.exists(f.name):
             os.remove(f.name)
-
-
-def named_apply(
-    fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False
-) -> nn.Module:
-    if not depth_first and include_root:
-        fn(module=module, name=name)
-    for child_name, child_module in module.named_children():
-        child_name = ".".join((name, child_name)) if name else child_name
-        named_apply(
-            fn=fn,
-            module=child_module,
-            name=child_name,
-            depth_first=depth_first,
-            include_root=True,
-        )
-    if depth_first and include_root:
-        fn(module=module, name=name)
-    return module
-
-
-def named_modules(module: nn.Module, name="", depth_first=True, include_root=False):
-    if not depth_first and include_root:
-        yield name, module
-    for child_name, child_module in module.named_children():
-        child_name = ".".join((name, child_name)) if name else child_name
-        yield from named_modules(
-            module=child_module,
-            name=child_name,
-            depth_first=depth_first,
-            include_root=True,
-        )
-    if depth_first and include_root:
-        yield name, module