ghlee94's picture
Init
2a13495
import torch
import torch.nn as nn
from typing import List
from collections import OrderedDict
from . import _utils as utils
class EncoderMixin:
"""Add encoder functionality such as:
- output channels specification of feature tensors (produced by encoder)
- patching first convolution for arbitrary input channels
"""
_output_stride = 32
@property
def out_channels(self):
"""Return channels dimensions for each tensor of forward output of encoder"""
return self._out_channels[: self._depth + 1]
@property
def output_stride(self):
return min(self._output_stride, 2 ** self._depth)
def set_in_channels(self, in_channels, pretrained=True):
"""Change first convolution channels"""
if in_channels == 3:
return
self._in_channels = in_channels
if self._out_channels[0] == 3:
self._out_channels = tuple([in_channels] + list(self._out_channels)[1:])
utils.patch_first_conv(
model=self, new_in_channels=in_channels, pretrained=pretrained
)
def get_stages(self):
"""Override it in your implementation"""
raise NotImplementedError
def make_dilated(self, output_stride):
if output_stride == 16:
stage_list = [
5,
]
dilation_list = [
2,
]
elif output_stride == 8:
stage_list = [4, 5]
dilation_list = [2, 4]
else:
raise ValueError(
"Output stride should be 16 or 8, got {}.".format(output_stride)
)
self._output_stride = output_stride
stages = self.get_stages()
for stage_indx, dilation_rate in zip(stage_list, dilation_list):
utils.replace_strides_with_dilation(
module=stages[stage_indx], dilation_rate=dilation_rate,
)