179 lines
6.2 KiB
Python
179 lines
6.2 KiB
Python
""" Each encoder should have following attributes and methods and be inherited from `_base.EncoderMixin`
|
|
|
|
Attributes:
|
|
|
|
_out_channels (list of int): specify number of channels for each encoder feature tensor
|
|
_depth (int): specify number of stages in decoder (in other words number of downsampling operations)
|
|
_in_channels (int): default number of input channels in first Conv2d layer for encoder (usually 3)
|
|
|
|
Methods:
|
|
|
|
forward(self, x: torch.Tensor)
|
|
produce list of features of different spatial resolutions, each feature is a 4D torch.tensor of
|
|
shape NCHW (features should be sorted in descending order according to spatial resolution, starting
|
|
with resolution same as input `x` tensor).
|
|
|
|
Input: `x` with shape (1, 3, 64, 64)
|
|
Output: [f0, f1, f2, f3, f4, f5] - features with corresponding shapes
|
|
[(1, 3, 64, 64), (1, 64, 32, 32), (1, 128, 16, 16), (1, 256, 8, 8),
|
|
(1, 512, 4, 4), (1, 1024, 2, 2)] (C - dim may differ)
|
|
|
|
also should support number of features according to specified depth, e.g. if depth = 5,
|
|
number of feature tensors = 6 (one with same resolution as input and 5 downsampled),
|
|
depth = 3 -> number of feature tensors = 4 (one with same resolution as input and 3 downsampled).
|
|
"""
|
|
import torch.nn as nn
|
|
from efficientnet_pytorch import EfficientNet
|
|
from efficientnet_pytorch.utils import url_map, url_map_advprop, get_model_params
|
|
|
|
from ._base import EncoderMixin
|
|
|
|
|
|
class EfficientNetEncoder(EfficientNet, EncoderMixin):
|
|
def __init__(self, stage_idxs, out_channels, model_name, depth=5):
|
|
|
|
blocks_args, global_params = get_model_params(model_name, override_params=None)
|
|
super().__init__(blocks_args, global_params)
|
|
|
|
self._stage_idxs = stage_idxs
|
|
self._out_channels = out_channels
|
|
self._depth = depth
|
|
self._in_channels = 3
|
|
|
|
del self._fc
|
|
|
|
def get_stages(self):
|
|
return [
|
|
nn.Identity(),
|
|
nn.Sequential(self._conv_stem, self._bn0, self._swish),
|
|
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()
|
|
|
|
block_number = 0.
|
|
drop_connect_rate = self._global_params.drop_connect_rate
|
|
|
|
features = []
|
|
for i in range(self._depth + 1):
|
|
|
|
# Identity and Sequential stages
|
|
if i < 2:
|
|
x = stages[i](x)
|
|
|
|
# Block stages need drop_connect rate
|
|
else:
|
|
for module in stages[i]:
|
|
drop_connect = drop_connect_rate * block_number / len(self._blocks)
|
|
block_number += 1.
|
|
x = module(x, drop_connect)
|
|
|
|
features.append(x)
|
|
|
|
return features
|
|
|
|
def load_state_dict(self, state_dict, **kwargs):
|
|
state_dict.pop("_fc.bias", None)
|
|
state_dict.pop("_fc.weight", None)
|
|
super().load_state_dict(state_dict, **kwargs)
|
|
|
|
|
|
def _get_pretrained_settings(encoder):
|
|
pretrained_settings = {
|
|
"imagenet": {
|
|
"mean": [0.485, 0.456, 0.406],
|
|
"std": [0.229, 0.224, 0.225],
|
|
"url": url_map[encoder],
|
|
"input_space": "RGB",
|
|
"input_range": [0, 1],
|
|
},
|
|
"advprop": {
|
|
"mean": [0.5, 0.5, 0.5],
|
|
"std": [0.5, 0.5, 0.5],
|
|
"url": url_map_advprop[encoder],
|
|
"input_space": "RGB",
|
|
"input_range": [0, 1],
|
|
}
|
|
}
|
|
return pretrained_settings
|
|
|
|
|
|
efficient_net_encoders = {
|
|
"efficientnet-b0": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b0"),
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 40, 112, 320),
|
|
"stage_idxs": (3, 5, 9, 16),
|
|
"model_name": "efficientnet-b0",
|
|
},
|
|
},
|
|
"efficientnet-b1": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b1"),
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 40, 112, 320),
|
|
"stage_idxs": (5, 8, 16, 23),
|
|
"model_name": "efficientnet-b1",
|
|
},
|
|
},
|
|
"efficientnet-b2": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b2"),
|
|
"params": {
|
|
"out_channels": (3, 32, 24, 48, 120, 352),
|
|
"stage_idxs": (5, 8, 16, 23),
|
|
"model_name": "efficientnet-b2",
|
|
},
|
|
},
|
|
"efficientnet-b3": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b3"),
|
|
"params": {
|
|
"out_channels": (3, 40, 32, 48, 136, 384),
|
|
"stage_idxs": (5, 8, 18, 26),
|
|
"model_name": "efficientnet-b3",
|
|
},
|
|
},
|
|
"efficientnet-b4": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b4"),
|
|
"params": {
|
|
"out_channels": (3, 48, 32, 56, 160, 448),
|
|
"stage_idxs": (6, 10, 22, 32),
|
|
"model_name": "efficientnet-b4",
|
|
},
|
|
},
|
|
"efficientnet-b5": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b5"),
|
|
"params": {
|
|
"out_channels": (3, 48, 40, 64, 176, 512),
|
|
"stage_idxs": (8, 13, 27, 39),
|
|
"model_name": "efficientnet-b5",
|
|
},
|
|
},
|
|
"efficientnet-b6": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b6"),
|
|
"params": {
|
|
"out_channels": (3, 56, 40, 72, 200, 576),
|
|
"stage_idxs": (9, 15, 31, 45),
|
|
"model_name": "efficientnet-b6",
|
|
},
|
|
},
|
|
"efficientnet-b7": {
|
|
"encoder": EfficientNetEncoder,
|
|
"pretrained_settings": _get_pretrained_settings("efficientnet-b7"),
|
|
"params": {
|
|
"out_channels": (3, 64, 48, 80, 224, 640),
|
|
"stage_idxs": (11, 18, 38, 55),
|
|
"model_name": "efficientnet-b7",
|
|
},
|
|
},
|
|
}
|