Features/edgevit3

docs/source/model_zoo_cls.md

@@ -82,8 +82,8 @@
 | efficientformer_l1 | [efficientformer_l1](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/efficientformer/efficientformer_l1.py) | 80.102   | 94.934   | 1820    | 7.5    | [model](https://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/efficientformer/efficientformer_l1_1000d.pth) |
 | efficientformer_l3 | [efficientformer_l3](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/efficientformer/efficientformer_l3.py) | 82.272   | 96.028   | 2436    | 13.07    | [model](https://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/efficientformer/efficientformer_l3_300d.pth) |
 | efficientformer_l7 | [efficientformer_l7](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/efficientformer/efficientformer_l7.py) | 83.076   | 96.44   | 1622    | 18.96    | [model](https://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/efficientformer/efficientformer_l7_300d.pth) |
-| EdgeVit_xxs_b512_224 | [EdgeVit_xxs_b512_224](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/edgevit/imagenet_edgeVIT_xxs_jpg.py) | 75.18   | 92.188   | 206   | 8.67    | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edgexxs/ClsEvaluator_neck_top1_best.pth) |
-| EdgeVit_xs_b256_224 | [EdgeVit_xs_b256_224](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/edgevit/imagenet_edgeVIT_xs_jpg.py) | 77.624   | 93.47  | 551   | 8.04    | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edgexs/ClsEvaluator_neck_top1_best.pth) |
-| EdgeVit_s_b128_224 | [EdgeVit_s_b128_224](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/edgevit/imagenet_edgeVIT_s_jpg.py) | 80.3   | 95.302  | 576   | 13.49   | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edges/ClsEvaluator_neck_top1_best.pth) |
+| EdgeVit_xxs_b512_224 | [EdgeVit_xxs_b512_224](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/edgevit/imagenet_edgeVIT_xxs_jpg.py) | 75.18   | 92.188   | 206   | 8.67    | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edgexxs/edgevit_xxs.pth) |
+| EdgeVit_xs_b256_224 | [EdgeVit_xs_b256_224](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/edgevit/imagenet_edgeVIT_xs_jpg.py) | 77.624   | 93.47  | 551   | 8.04    | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edgexs/edgevit_xs.pth) |
+| EdgeVit_s_b128_224 | [EdgeVit_s_b128_224](https://github.com/alibaba/EasyCV/tree/master/configs/classification/imagenet/edgevit/imagenet_edgeVIT_s_jpg.py) | 80.3   | 95.302  | 576   | 13.49   | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edges/edgevit_s.pth) |
 
 (ps: 通过导入官方模型得到推理结果，需要torch.__version__ >= 1.9.0，推理的输入尺寸默认为224，机器默认为V100 16G，其中gpu memory记录的是gpu peak memory)

easycv/models/backbones/edgevit.py

@@ -14,6 +14,7 @@
 from easycv.models.utils import ConvMlp, Mlp
 from easycv.utils.checkpoint import load_checkpoint
 from easycv.utils.logger import get_root_logger
+from ..modelzoo import EdgeVit as model_urls
 from ..registry import BACKBONES
 
 
@@ -249,8 +250,7 @@ def __init__(self,
                  attn_drop_rate=0.,
                  drop_path_rate=0.,
                  norm_layer=partial(nn.LayerNorm, eps=1e-8),
-                 sr_ratios=[4, 2, 2, 1],
-                 pretrained=None):
+                 sr_ratios=[4, 2, 2, 1]):
         """
         Args:
             depth (list): depth of each stage
@@ -362,16 +362,21 @@ def __init__(self,
         else:
             self.pre_logits = nn.Identity()
 
-        self.pretrained = pretrained
         self.init_weights()
 
-    def init_weights(self, pretrained=None):
+        size_dict = {
+            'xxs': [1, 1, 3, 2],
+            'xs': [1, 1, 3, 1],
+            's': [1, 2, 5, 3]
+        }
+
+        self.default_pretrained_model_path = model_urls.get(
+            self.__class__.__name__ + '_' +
+            [k for k, v in size_dict.items() if v == depth][0], None)
+
+    def init_weights(self):
         """Initialize the weights in backbone.
-        Args:
-            pretrained (str, optional): Path to pre-trained weights.
-                Defaults to None.
         """
-        pretrained = pretrained or self.pretrained
 
         def _init_weights(m):
             if isinstance(m, nn.Linear):
@@ -382,15 +387,6 @@ def _init_weights(m):
                 nn.init.constant_(m.bias, 0)
                 nn.init.constant_(m.weight, 1.0)
 
-        if isinstance(pretrained, str):
-            self.apply(_init_weights)
-            logger = get_root_logger()
-            load_checkpoint(self, pretrained, strict=False, logger=logger)
-        elif pretrained is None:
-            self.apply(_init_weights)
-        else:
-            raise TypeError('pretrained must be a str or None')
-
     @torch.jit.ignore
     def no_weight_decay(self):
         return {'pos_embed', 'cls_token'}

easycv/models/modelzoo.py

@@ -279,3 +279,12 @@
     'dynamic_vit_large_p16':
     'https://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/timm/vit/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npz',
 }
+
+EdgeVit = {
+    'EdgeVit_xxs':
+    'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edgexxs/edgevit_xxs.pth',
+    'EdgeVit_xs':
+    'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edgexs/edgevit_xs.pth',
+    'EdgeVit_s':
+    'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/edgevit/edges/edgevit_s.pth',
+}