File size: 1,927 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 |
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,
)
|