Spaces:
Runtime error
Runtime error
File size: 1,684 Bytes
5e85cf4 |
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 |
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
"""
@property
def out_channels(self):
"""Return channels dimensions for each tensor of forward output of encoder"""
return self._out_channels[: self._depth + 1]
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):
"""Method should be overridden in encoder"""
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))
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,
)
|