# Copyright (c) Alibaba, Inc. and its affiliates.
import importlib
import timm
import torch
import torch.nn as nn
from timm.models.helpers import load_pretrained
from timm.models.hub import download_cached_file
from easycv.framework.errors import ValueError
from easycv.utils.logger import get_root_logger, print_log
from ..modelzoo import timm_models as model_urls
from ..registry import BACKBONES
from .shuffle_transformer import (shuffletrans_base_p4_w7_224,
shuffletrans_small_p4_w7_224,
shuffletrans_tiny_p4_w7_224)
from .swin_transformer_dynamic import (dynamic_swin_base_p4_w7_224,
dynamic_swin_small_p4_w7_224,
dynamic_swin_tiny_p4_w7_224)
from .vit_transformer_dynamic import (dynamic_deit_small_p16,
dynamic_deit_tiny_p16,
dynamic_vit_base_p16,
dynamic_vit_huge_p14,
dynamic_vit_large_p16)
from .xcit_transformer import (xcit_large_24_p8, xcit_medium_24_p8,
xcit_medium_24_p16, xcit_small_12_p8,
xcit_small_12_p16)
_MODEL_MAP = {
# shuffle_transformer
'shuffletrans_tiny_p4_w7_224': shuffletrans_tiny_p4_w7_224,
'shuffletrans_base_p4_w7_224': shuffletrans_base_p4_w7_224,
'shuffletrans_small_p4_w7_224': shuffletrans_small_p4_w7_224,
# swin_transformer_dynamic
'dynamic_swin_tiny_p4_w7_224': dynamic_swin_tiny_p4_w7_224,
'dynamic_swin_small_p4_w7_224': dynamic_swin_small_p4_w7_224,
'dynamic_swin_base_p4_w7_224': dynamic_swin_base_p4_w7_224,
# vit_transformer_dynamic
'dynamic_deit_small_p16': dynamic_deit_small_p16,
'dynamic_deit_tiny_p16': dynamic_deit_tiny_p16,
'dynamic_vit_base_p16': dynamic_vit_base_p16,
'dynamic_vit_large_p16': dynamic_vit_large_p16,
'dynamic_vit_huge_p14': dynamic_vit_huge_p14,
# xcit_transformer
'xcit_small_12_p16': xcit_small_12_p16,
'xcit_small_12_p8': xcit_small_12_p8,
'xcit_medium_24_p16': xcit_medium_24_p16,
'xcit_medium_24_p8': xcit_medium_24_p8,
'xcit_large_24_p8': xcit_large_24_p8
}
@BACKBONES.register_module
class PytorchImageModelWrapper(nn.Module):
"""Support Backbones From pytorch-image-models.
The PyTorch community has lots of awesome contributions for image models. PyTorch Image Models (timm) is
a collection of image models, aim to pull together a wide variety of SOTA models with ability to reproduce
ImageNet training results.
Model pages can be found at https://rwightman.github.io/pytorch-image-models/models/
References: https://github.com/rwightman/pytorch-image-models
"""
def __init__(self,
model_name='resnet50',
scriptable=None,
exportable=None,
no_jit=None,
**kwargs):
"""
Inits PytorchImageModelWrapper by timm.create_models
Args:
model_name (str): name of model to instantiate
scriptable (bool): set layer config so that model is jit scriptable (not working for all models yet)
exportable (bool): set layer config so that model is traceable / ONNX exportable (not fully impl/obeyed yet)
no_jit (bool): set layer config so that model doesn't utilize jit scripted layers (so far activations only)
"""
super(PytorchImageModelWrapper, self).__init__()
self.model_name = model_name
timm_model_names = timm.list_models(pretrained=False)
self.timm_model_names = timm_model_names
assert model_name in timm_model_names or model_name in _MODEL_MAP, \
f'{model_name} is not in model_list of timm/fair, please check the model_name!'
# Default to use backbone without head from timm
if 'num_classes' not in kwargs:
kwargs['num_classes'] = 0
self.num_classes = kwargs['num_classes']
# create model by timm
if model_name in timm_model_names:
self.model = timm.create_model(model_name, False, '', scriptable,
exportable, no_jit, **kwargs)
elif model_name in _MODEL_MAP:
self.model = _MODEL_MAP[model_name](**kwargs)
def init_weights(self, pretrained=None):
"""
Args:
if pretrained == True, load model from default path;
if pretrained == False or None, load from init weights.
if model_name in timm_model_names, load model from timm default path;
if model_name in _MODEL_MAP, load model from easycv default path
"""
logger = get_root_logger()
if pretrained:
if self.model_name in self.timm_model_names:
if self.model_name in model_urls:
default_pretrained_model_path = model_urls[self.model_name]
print_log(
'load model from default path: {}'.format(
default_pretrained_model_path), logger)
if default_pretrained_model_path.endswith('.npz'):
pretrained_loc = download_cached_file(
default_pretrained_model_path,
check_hash=False,
progress=False)
return self.model.load_pretrained(pretrained_loc)
else:
backbone_module = importlib.import_module(
self.model.__module__)
return load_pretrained(
self.model,
default_cfg={
'url': default_pretrained_model_path,
'classifier': 'head',
'num_classes': 1000
},
num_classes=self.num_classes,
filter_fn=backbone_module.checkpoint_filter_fn
if hasattr(backbone_module, 'checkpoint_filter_fn')
else None,
strict=False)
else:
logger.warning('pretrained model for model_name not found')
elif self.model_name in _MODEL_MAP:
if self.model_name in model_urls.keys():
default_pretrained_model_path = model_urls[self.model_name]
print_log(
'load model from default path: {}'.format(
default_pretrained_model_path), logger)
try_max = 3
try_idx = 0
while try_idx < try_max:
try:
state_dict = torch.hub.load_state_dict_from_url(
url=default_pretrained_model_path,
map_location='cpu',
)
try_idx += try_max
except Exception:
try_idx += 1
state_dict = {}
if try_idx == try_max:
print_log(
f'load from url failed ! oh my DLC & OSS, you boys really good! {model_urls[self.model_name]}',
logger)
if 'model' in state_dict:
state_dict = state_dict['model']
self.model.load_state_dict(state_dict, strict=False)
else:
raise ValueError('{} not in evtorch modelzoo!'.format(
self.model_name))
else:
raise ValueError(
'Error: Fail to create {} with (pretrained={}...)'.format(
self.model_name, pretrained))
else:
print_log('load model from init weights')
def forward(self, x):
o = self.model(x)
if type(o) == tuple or type(o) == list:
features = []
for feature in o:
while feature.dim() < 4:
feature = feature.unsqueeze(-1)
features.append(feature)
else:
while o.dim() < 4:
o = o.unsqueeze(-1)
features = [o]
return tuple(features)