File size: 5,308 Bytes
2a13495 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
"""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 re
import torch.nn as nn
from pretrainedmodels.models.torchvision_models import pretrained_settings
from torchvision.models.densenet import DenseNet
from ._base import EncoderMixin
class TransitionWithSkip(nn.Module):
def __init__(self, module):
super().__init__()
self.module = module
def forward(self, x):
for module in self.module:
x = module(x)
if isinstance(module, nn.ReLU):
skip = x
return x, skip
class DenseNetEncoder(DenseNet, EncoderMixin):
def __init__(self, out_channels, depth=5, **kwargs):
super().__init__(**kwargs)
self._out_channels = out_channels
self._depth = depth
self._in_channels = 3
del self.classifier
def make_dilated(self, *args, **kwargs):
raise ValueError(
"DenseNet encoders do not support dilated mode "
"due to pooling operation for downsampling!"
)
def get_stages(self):
return [
nn.Identity(),
nn.Sequential(
self.features.conv0, self.features.norm0, self.features.relu0
),
nn.Sequential(
self.features.pool0,
self.features.denseblock1,
TransitionWithSkip(self.features.transition1),
),
nn.Sequential(
self.features.denseblock2, TransitionWithSkip(self.features.transition2)
),
nn.Sequential(
self.features.denseblock3, TransitionWithSkip(self.features.transition3)
),
nn.Sequential(self.features.denseblock4, self.features.norm5),
]
def forward(self, x):
stages = self.get_stages()
features = []
for i in range(self._depth + 1):
x = stages[i](x)
if isinstance(x, (list, tuple)):
x, skip = x
features.append(skip)
else:
features.append(x)
return features
def load_state_dict(self, state_dict):
pattern = re.compile(
r"^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$"
)
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
# remove linear
state_dict.pop("classifier.bias", None)
state_dict.pop("classifier.weight", None)
super().load_state_dict(state_dict)
densenet_encoders = {
"densenet121": {
"encoder": DenseNetEncoder,
"pretrained_settings": pretrained_settings["densenet121"],
"params": {
"out_channels": (3, 64, 256, 512, 1024, 1024),
"num_init_features": 64,
"growth_rate": 32,
"block_config": (6, 12, 24, 16),
},
},
"densenet169": {
"encoder": DenseNetEncoder,
"pretrained_settings": pretrained_settings["densenet169"],
"params": {
"out_channels": (3, 64, 256, 512, 1280, 1664),
"num_init_features": 64,
"growth_rate": 32,
"block_config": (6, 12, 32, 32),
},
},
"densenet201": {
"encoder": DenseNetEncoder,
"pretrained_settings": pretrained_settings["densenet201"],
"params": {
"out_channels": (3, 64, 256, 512, 1792, 1920),
"num_init_features": 64,
"growth_rate": 32,
"block_config": (6, 12, 48, 32),
},
},
"densenet161": {
"encoder": DenseNetEncoder,
"pretrained_settings": pretrained_settings["densenet161"],
"params": {
"out_channels": (3, 96, 384, 768, 2112, 2208),
"num_init_features": 96,
"growth_rate": 48,
"block_config": (6, 12, 36, 24),
},
},
}
|