176 lines
6.0 KiB
Python
176 lines
6.0 KiB
Python
import timm
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
|
|
from ._base import EncoderMixin
|
|
|
|
|
|
def _make_divisible(x, divisible_by=8):
|
|
return int(np.ceil(x * 1. / divisible_by) * divisible_by)
|
|
|
|
|
|
class MobileNetV3Encoder(nn.Module, EncoderMixin):
|
|
def __init__(self, model_name, width_mult, depth=5, **kwargs):
|
|
super().__init__()
|
|
if "large" not in model_name and "small" not in model_name:
|
|
raise ValueError(
|
|
'MobileNetV3 wrong model name {}'.format(model_name)
|
|
)
|
|
|
|
self._mode = "small" if "small" in model_name else "large"
|
|
self._depth = depth
|
|
self._out_channels = self._get_channels(self._mode, width_mult)
|
|
self._in_channels = 3
|
|
|
|
# minimal models replace hardswish with relu
|
|
self.model = timm.create_model(
|
|
model_name=model_name,
|
|
scriptable=True, # torch.jit scriptable
|
|
exportable=True, # onnx export
|
|
features_only=True,
|
|
)
|
|
|
|
def _get_channels(self, mode, width_mult):
|
|
if mode == "small":
|
|
channels = [16, 16, 24, 48, 576]
|
|
else:
|
|
channels = [16, 24, 40, 112, 960]
|
|
channels = [3,] + [_make_divisible(x * width_mult) for x in channels]
|
|
return tuple(channels)
|
|
|
|
def get_stages(self):
|
|
if self._mode == 'small':
|
|
return [
|
|
nn.Identity(),
|
|
nn.Sequential(
|
|
self.model.conv_stem,
|
|
self.model.bn1,
|
|
self.model.act1,
|
|
),
|
|
self.model.blocks[0],
|
|
self.model.blocks[1],
|
|
self.model.blocks[2:4],
|
|
self.model.blocks[4:],
|
|
]
|
|
elif self._mode == 'large':
|
|
return [
|
|
nn.Identity(),
|
|
nn.Sequential(
|
|
self.model.conv_stem,
|
|
self.model.bn1,
|
|
self.model.act1,
|
|
self.model.blocks[0],
|
|
),
|
|
self.model.blocks[1],
|
|
self.model.blocks[2],
|
|
self.model.blocks[3:5],
|
|
self.model.blocks[5:],
|
|
]
|
|
else:
|
|
ValueError('MobileNetV3 mode should be small or large, got {}'.format(self._mode))
|
|
|
|
def forward(self, x):
|
|
stages = self.get_stages()
|
|
|
|
features = []
|
|
for i in range(self._depth + 1):
|
|
x = stages[i](x)
|
|
features.append(x)
|
|
|
|
return features
|
|
|
|
def load_state_dict(self, state_dict, **kwargs):
|
|
state_dict.pop('conv_head.weight', None)
|
|
state_dict.pop('conv_head.bias', None)
|
|
state_dict.pop('classifier.weight', None)
|
|
state_dict.pop('classifier.bias', None)
|
|
self.model.load_state_dict(state_dict, **kwargs)
|
|
|
|
|
|
mobilenetv3_weights = {
|
|
'tf_mobilenetv3_large_075': {
|
|
'imagenet': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_075-150ee8b0.pth'
|
|
},
|
|
'tf_mobilenetv3_large_100': {
|
|
'imagenet': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_100-427764d5.pth'
|
|
},
|
|
'tf_mobilenetv3_large_minimal_100': {
|
|
'imagenet': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_minimal_100-8596ae28.pth'
|
|
},
|
|
'tf_mobilenetv3_small_075': {
|
|
'imagenet': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_075-da427f52.pth'
|
|
},
|
|
'tf_mobilenetv3_small_100': {
|
|
'imagenet': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_100-37f49e2b.pth'
|
|
},
|
|
'tf_mobilenetv3_small_minimal_100': {
|
|
'imagenet': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_minimal_100-922a7843.pth'
|
|
},
|
|
|
|
|
|
}
|
|
|
|
pretrained_settings = {}
|
|
for model_name, sources in mobilenetv3_weights.items():
|
|
pretrained_settings[model_name] = {}
|
|
for source_name, source_url in sources.items():
|
|
pretrained_settings[model_name][source_name] = {
|
|
"url": source_url,
|
|
'input_range': [0, 1],
|
|
'mean': [0.485, 0.456, 0.406],
|
|
'std': [0.229, 0.224, 0.225],
|
|
'input_space': 'RGB',
|
|
}
|
|
|
|
|
|
timm_mobilenetv3_encoders = {
|
|
'timm-mobilenetv3_large_075': {
|
|
'encoder': MobileNetV3Encoder,
|
|
'pretrained_settings': pretrained_settings['tf_mobilenetv3_large_075'],
|
|
'params': {
|
|
'model_name': 'tf_mobilenetv3_large_075',
|
|
'width_mult': 0.75
|
|
}
|
|
},
|
|
'timm-mobilenetv3_large_100': {
|
|
'encoder': MobileNetV3Encoder,
|
|
'pretrained_settings': pretrained_settings['tf_mobilenetv3_large_100'],
|
|
'params': {
|
|
'model_name': 'tf_mobilenetv3_large_100',
|
|
'width_mult': 1.0
|
|
}
|
|
},
|
|
'timm-mobilenetv3_large_minimal_100': {
|
|
'encoder': MobileNetV3Encoder,
|
|
'pretrained_settings': pretrained_settings['tf_mobilenetv3_large_minimal_100'],
|
|
'params': {
|
|
'model_name': 'tf_mobilenetv3_large_minimal_100',
|
|
'width_mult': 1.0
|
|
}
|
|
},
|
|
'timm-mobilenetv3_small_075': {
|
|
'encoder': MobileNetV3Encoder,
|
|
'pretrained_settings': pretrained_settings['tf_mobilenetv3_small_075'],
|
|
'params': {
|
|
'model_name': 'tf_mobilenetv3_small_075',
|
|
'width_mult': 0.75
|
|
}
|
|
},
|
|
'timm-mobilenetv3_small_100': {
|
|
'encoder': MobileNetV3Encoder,
|
|
'pretrained_settings': pretrained_settings['tf_mobilenetv3_small_100'],
|
|
'params': {
|
|
'model_name': 'tf_mobilenetv3_small_100',
|
|
'width_mult': 1.0
|
|
}
|
|
},
|
|
'timm-mobilenetv3_small_minimal_100': {
|
|
'encoder': MobileNetV3Encoder,
|
|
'pretrained_settings': pretrained_settings['tf_mobilenetv3_small_minimal_100'],
|
|
'params': {
|
|
'model_name': 'tf_mobilenetv3_small_minimal_100',
|
|
'width_mult': 1.0
|
|
}
|
|
},
|
|
}
|