|
"""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.0 |
|
drop_connect_rate = self._global_params.drop_connect_rate |
|
|
|
features = [] |
|
for i in range(self._depth + 1): |
|
|
|
|
|
if i < 2: |
|
x = stages[i](x) |
|
|
|
|
|
else: |
|
for module in stages[i]: |
|
drop_connect = drop_connect_rate * block_number / len(self._blocks) |
|
block_number += 1.0 |
|
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", |
|
}, |
|
}, |
|
} |
|
|
