diff --git a/docs/source/model_zoo_cls.md b/docs/source/model_zoo_cls.md index 521f527b..6171a2a1 100644 --- a/docs/source/model_zoo_cls.md +++ b/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) diff --git a/easycv/models/backbones/edgevit.py b/easycv/models/backbones/edgevit.py index f2f4499e..405d9d3c 100644 --- a/easycv/models/backbones/edgevit.py +++ b/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'} diff --git a/easycv/models/modelzoo.py b/easycv/models/modelzoo.py index d367cc62..423484c0 100644 --- a/easycv/models/modelzoo.py +++ b/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', +}