383 lines
13 KiB
Python
383 lines
13 KiB
Python
from functools import partial
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from timm.models.efficientnet import EfficientNet
|
|
from timm.models.efficientnet import decode_arch_def, round_channels, default_cfgs
|
|
from timm.models.layers.activations import Swish
|
|
|
|
from ._base import EncoderMixin
|
|
|
|
|
|
def get_efficientnet_kwargs(channel_multiplier=1.0, depth_multiplier=1.0, drop_rate=0.2):
|
|
"""Creates an EfficientNet model.
|
|
Ref impl: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py
|
|
Paper: https://arxiv.org/abs/1905.11946
|
|
EfficientNet params
|
|
name: (channel_multiplier, depth_multiplier, resolution, dropout_rate)
|
|
'efficientnet-b0': (1.0, 1.0, 224, 0.2),
|
|
'efficientnet-b1': (1.0, 1.1, 240, 0.2),
|
|
'efficientnet-b2': (1.1, 1.2, 260, 0.3),
|
|
'efficientnet-b3': (1.2, 1.4, 300, 0.3),
|
|
'efficientnet-b4': (1.4, 1.8, 380, 0.4),
|
|
'efficientnet-b5': (1.6, 2.2, 456, 0.4),
|
|
'efficientnet-b6': (1.8, 2.6, 528, 0.5),
|
|
'efficientnet-b7': (2.0, 3.1, 600, 0.5),
|
|
'efficientnet-b8': (2.2, 3.6, 672, 0.5),
|
|
'efficientnet-l2': (4.3, 5.3, 800, 0.5),
|
|
Args:
|
|
channel_multiplier: multiplier to number of channels per layer
|
|
depth_multiplier: multiplier to number of repeats per stage
|
|
"""
|
|
arch_def = [
|
|
['ds_r1_k3_s1_e1_c16_se0.25'],
|
|
['ir_r2_k3_s2_e6_c24_se0.25'],
|
|
['ir_r2_k5_s2_e6_c40_se0.25'],
|
|
['ir_r3_k3_s2_e6_c80_se0.25'],
|
|
['ir_r3_k5_s1_e6_c112_se0.25'],
|
|
['ir_r4_k5_s2_e6_c192_se0.25'],
|
|
['ir_r1_k3_s1_e6_c320_se0.25'],
|
|
]
|
|
model_kwargs = dict(
|
|
block_args=decode_arch_def(arch_def, depth_multiplier),
|
|
num_features=round_channels(1280, channel_multiplier, 8, None),
|
|
stem_size=32,
|
|
round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
|
|
act_layer=Swish,
|
|
drop_rate=drop_rate,
|
|
drop_path_rate=0.2,
|
|
)
|
|
return model_kwargs
|
|
|
|
def gen_efficientnet_lite_kwargs(channel_multiplier=1.0, depth_multiplier=1.0, drop_rate=0.2):
|
|
"""Creates an EfficientNet-Lite model.
|
|
|
|
Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite
|
|
Paper: https://arxiv.org/abs/1905.11946
|
|
|
|
EfficientNet params
|
|
name: (channel_multiplier, depth_multiplier, resolution, dropout_rate)
|
|
'efficientnet-lite0': (1.0, 1.0, 224, 0.2),
|
|
'efficientnet-lite1': (1.0, 1.1, 240, 0.2),
|
|
'efficientnet-lite2': (1.1, 1.2, 260, 0.3),
|
|
'efficientnet-lite3': (1.2, 1.4, 280, 0.3),
|
|
'efficientnet-lite4': (1.4, 1.8, 300, 0.3),
|
|
|
|
Args:
|
|
channel_multiplier: multiplier to number of channels per layer
|
|
depth_multiplier: multiplier to number of repeats per stage
|
|
"""
|
|
arch_def = [
|
|
['ds_r1_k3_s1_e1_c16'],
|
|
['ir_r2_k3_s2_e6_c24'],
|
|
['ir_r2_k5_s2_e6_c40'],
|
|
['ir_r3_k3_s2_e6_c80'],
|
|
['ir_r3_k5_s1_e6_c112'],
|
|
['ir_r4_k5_s2_e6_c192'],
|
|
['ir_r1_k3_s1_e6_c320'],
|
|
]
|
|
model_kwargs = dict(
|
|
block_args=decode_arch_def(arch_def, depth_multiplier, fix_first_last=True),
|
|
num_features=1280,
|
|
stem_size=32,
|
|
fix_stem=True,
|
|
round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
|
|
act_layer=nn.ReLU6,
|
|
drop_rate=drop_rate,
|
|
drop_path_rate=0.2,
|
|
)
|
|
return model_kwargs
|
|
|
|
class EfficientNetBaseEncoder(EfficientNet, EncoderMixin):
|
|
|
|
def __init__(self, stage_idxs, out_channels, depth=5, **kwargs):
|
|
super().__init__(**kwargs)
|
|
|
|
self._stage_idxs = stage_idxs
|
|
self._out_channels = out_channels
|
|
self._depth = depth
|
|
self._in_channels = 3
|
|
|
|
del self.classifier
|
|
|
|
def get_stages(self):
|
|
return [
|
|
nn.Identity(),
|
|
nn.Sequential(self.conv_stem, self.bn1, self.act1),
|
|
self.blocks[:self._stage_idxs[0]],
|
|
self.blocks[self._stage_idxs[0]:self._stage_idxs[1]],
|
|
self.blocks[self._stage_idxs[1]:self._stage_idxs[2]],
|
|
self.blocks[self._stage_idxs[2]:],
|
|
]
|
|
|
|
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("classifier.bias", None)
|
|
state_dict.pop("classifier.weight", None)
|
|
super().load_state_dict(state_dict, **kwargs)
|
|
|
|
|
|
class EfficientNetEncoder(EfficientNetBaseEncoder):
|
|
|
|
def __init__(self, stage_idxs, out_channels, depth=5, channel_multiplier=1.0, depth_multiplier=1.0, drop_rate=0.2):
|
|
kwargs = get_efficientnet_kwargs(channel_multiplier, depth_multiplier, drop_rate)
|
|
super().__init__(stage_idxs, out_channels, depth, **kwargs)
|
|
|
|
|
|
class EfficientNetLiteEncoder(EfficientNetBaseEncoder):
|
|
|
|
def __init__(self, stage_idxs, out_channels, depth=5, channel_multiplier=1.0, depth_multiplier=1.0, drop_rate=0.2):
|
|
kwargs = gen_efficientnet_lite_kwargs(channel_multiplier, depth_multiplier, drop_rate)
|
|
super().__init__(stage_idxs, out_channels, depth, **kwargs)
|
|
|
|
|
|
def prepare_settings(settings):
|
|
return {
|
|
"mean": settings["mean"],
|
|
"std": settings["std"],
|
|
"url": settings["url"],
|
|
"input_range": (0, 1),
|
|
"input_space": "RGB",
|
|
}
|
|
|
|
|
|
timm_efficientnet_encoders = {
|
|
|
|
"timm-efficientnet-b0": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b0"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b0_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b0_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 40, 112, 320),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.0,
|
|
"depth_multiplier": 1.0,
|
|
"drop_rate": 0.2,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b1": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b1"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b1_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b1_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 40, 112, 320),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.0,
|
|
"depth_multiplier": 1.1,
|
|
"drop_rate": 0.2,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b2": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b2"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b2_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b2_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 48, 120, 352),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.1,
|
|
"depth_multiplier": 1.2,
|
|
"drop_rate": 0.3,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b3": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b3"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b3_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b3_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 40, 32, 48, 136, 384),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.2,
|
|
"depth_multiplier": 1.4,
|
|
"drop_rate": 0.3,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b4": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b4"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b4_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b4_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 48, 32, 56, 160, 448),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.4,
|
|
"depth_multiplier": 1.8,
|
|
"drop_rate": 0.4,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b5": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b5"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b5_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b5_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 48, 40, 64, 176, 512),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.6,
|
|
"depth_multiplier": 2.2,
|
|
"drop_rate": 0.4,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b6": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b6"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b6_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b6_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 56, 40, 72, 200, 576),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.8,
|
|
"depth_multiplier": 2.6,
|
|
"drop_rate": 0.5,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b7": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b7"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b7_ap"]),
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_b7_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 64, 48, 80, 224, 640),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 2.0,
|
|
"depth_multiplier": 3.1,
|
|
"drop_rate": 0.5,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-b8": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_b8"]),
|
|
"advprop": prepare_settings(default_cfgs["tf_efficientnet_b8_ap"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 72, 56, 88, 248, 704),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 2.2,
|
|
"depth_multiplier": 3.6,
|
|
"drop_rate": 0.5,
|
|
},
|
|
},
|
|
|
|
"timm-efficientnet-l2": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": {
|
|
"noisy-student": prepare_settings(default_cfgs["tf_efficientnet_l2_ns"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 136, 104, 176, 480, 1376),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 4.3,
|
|
"depth_multiplier": 5.3,
|
|
"drop_rate": 0.5,
|
|
},
|
|
},
|
|
|
|
"timm-tf_efficientnet_lite0": {
|
|
"encoder": EfficientNetLiteEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_lite0"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 40, 112, 320),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.0,
|
|
"depth_multiplier": 1.0,
|
|
"drop_rate": 0.2,
|
|
},
|
|
},
|
|
|
|
"timm-tf_efficientnet_lite1": {
|
|
"encoder": EfficientNetLiteEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_lite1"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 40, 112, 320),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.0,
|
|
"depth_multiplier": 1.1,
|
|
"drop_rate": 0.2,
|
|
},
|
|
},
|
|
|
|
"timm-tf_efficientnet_lite2": {
|
|
"encoder": EfficientNetLiteEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_lite2"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 48, 120, 352),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.1,
|
|
"depth_multiplier": 1.2,
|
|
"drop_rate": 0.3,
|
|
},
|
|
},
|
|
|
|
"timm-tf_efficientnet_lite3": {
|
|
"encoder": EfficientNetLiteEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_lite3"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 32, 48, 136, 384),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.2,
|
|
"depth_multiplier": 1.4,
|
|
"drop_rate": 0.3,
|
|
},
|
|
},
|
|
|
|
"timm-tf_efficientnet_lite4": {
|
|
"encoder": EfficientNetLiteEncoder,
|
|
"pretrained_settings": {
|
|
"imagenet": prepare_settings(default_cfgs["tf_efficientnet_lite4"]),
|
|
},
|
|
"params": {
|
|
"out_channels": (3, 32, 32, 56, 160, 448),
|
|
"stage_idxs": (2, 3, 5),
|
|
"channel_multiplier": 1.4,
|
|
"depth_multiplier": 1.8,
|
|
"drop_rate": 0.4,
|
|
},
|
|
},
|
|
}
|