67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
import re
|
|
import torch.nn as nn
|
|
|
|
from pretrainedmodels.models.xception import pretrained_settings
|
|
from pretrainedmodels.models.xception import Xception
|
|
|
|
from ._base import EncoderMixin
|
|
|
|
|
|
class XceptionEncoder(Xception, EncoderMixin):
|
|
|
|
def __init__(self, out_channels, *args, depth=5, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self._out_channels = out_channels
|
|
self._depth = depth
|
|
self._in_channels = 3
|
|
|
|
# modify padding to maintain output shape
|
|
self.conv1.padding = (1, 1)
|
|
self.conv2.padding = (1, 1)
|
|
|
|
del self.fc
|
|
|
|
def make_dilated(self, output_stride):
|
|
raise ValueError("Xception encoder does not support dilated mode "
|
|
"due to pooling operation for downsampling!")
|
|
|
|
def get_stages(self):
|
|
return [
|
|
nn.Identity(),
|
|
nn.Sequential(self.conv1, self.bn1, self.relu, self.conv2, self.bn2, self.relu),
|
|
self.block1,
|
|
self.block2,
|
|
nn.Sequential(self.block3, self.block4, self.block5, self.block6, self.block7,
|
|
self.block8, self.block9, self.block10, self.block11),
|
|
nn.Sequential(self.block12, self.conv3, self.bn3, self.relu, self.conv4, self.bn4),
|
|
]
|
|
|
|
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):
|
|
# remove linear
|
|
state_dict.pop('fc.bias', None)
|
|
state_dict.pop('fc.weight', None)
|
|
|
|
super().load_state_dict(state_dict)
|
|
|
|
|
|
xception_encoders = {
|
|
'xception': {
|
|
'encoder': XceptionEncoder,
|
|
'pretrained_settings': pretrained_settings['xception'],
|
|
'params': {
|
|
'out_channels': (3, 64, 128, 256, 728, 2048),
|
|
}
|
|
},
|
|
}
|