def __init__(
self,
in_channels: int,
out_channels: int,
pool_size: int,
use_batchnorm: bool = True,
interpolation_mode: str = "bilinear",
align_corners: bool = True,
complexity: int = 0,
):
"""@TODO: Docs. Contribution is welcome."""
super().__init__()
self.interpolation_mode = interpolation_mode
self.align_corners = align_corners
if pool_size == 1:
use_batchnorm = False
self._block = nn.Sequential(
nn.AdaptiveAvgPool2d(output_size=(pool_size, pool_size)),
_get_block(
in_channels,
out_channels,
abn_block=partial(ABN, use_batchnorm=use_batchnorm),
complexity=complexity,
),
)
def __init__(self,
in_channels: int,
out_channels: int,
in_strides: int = None,
abn_block: nn.Module = ABN,
activation: str = "ReLU",
first_stride: int = 1,
second_stride: int = 1,
pool_first: bool = False,
upsample_scale: int = 2,
interpolation_mode: str = "nearest",
align_corners: bool = None,
**kwargs):
"""@TODO: Docs. Contribution is welcome."""
super().__init__(in_channels, out_channels, in_strides)
if in_strides is None:
self._out_strides = None
elif pool_first:
self._out_strides = in_strides * first_stride * second_stride * 2 // upsample_scale
else:
self._out_strides = in_strides * first_stride * second_stride // upsample_scale
self.pool_first = pool_first
self.upsample_scale = upsample_scale
self.interpolation_mode = interpolation_mode
self.align_corners = align_corners
self._block = _get_block(in_channels=in_channels,
out_channels=out_channels,
abn_block=abn_block,
activation=activation,
first_stride=first_stride,
second_stride=second_stride,
**kwargs)
def __init__(
self,
in_channels: List[int],
in_strides: List[int],
downsample_factor: int = 8,
use_batchnorm: bool = True,
out_channels: int = 512,
):
"""@TODO: Docs. Contribution is welcome."""
super().__init__(in_channels, in_strides)
self.block_offset = self._get_block_offset(downsample_factor)
psp_out_channels: int = self._get(in_channels)
self.psp = PSPBlock(
psp_out_channels, pool_sizes=(1, 2, 3, 6), use_batchnorm=use_batchnorm,
)
self.conv = _get_block(
psp_out_channels * 2,
out_channels,
kernel_size=1,
padding=0,
abn_block=partial(ABN, use_batchnorm=use_batchnorm),
complexity=0,
)
self._out_channels = out_channels
self.downsample_factor = downsample_factor
def __init__(self,
in_channels: int,
out_channels: int,
in_strides: int = None,
abn_block: nn.Module = ABN,
activation: str = "ReLU",
first_stride: int = 2,
second_stride: int = 1,
**kwargs):
"""@TODO: Docs. Contribution is welcome."""
super().__init__(in_channels, out_channels, in_strides)
self._out_strides = (in_strides * first_stride *
second_stride if in_strides is not None else None)
self._block = _get_block(in_channels=in_channels,
out_channels=out_channels,
abn_block=abn_block,
activation=activation,
first_stride=first_stride,
second_stride=second_stride,
**kwargs)
def _get_block(self,
abn_block: nn.Module = ABN,
activation: str = "ReLU",
pre_dropout_rate: float = 0.0,
post_dropout_rate: float = 0.0,
**kwargs):
layers = []
if pre_dropout_rate > 0:
layers.append(nn.Dropout2d(pre_dropout_rate, inplace=True))
layers.append(
_get_block(in_channels=self.in_channels + self.enc_channels,
out_channels=self.out_channels,
abn_block=abn_block,
activation=activation,
first_stride=1,
second_stride=1,
**kwargs))
if post_dropout_rate > 0:
layers.append(nn.Dropout2d(pre_dropout_rate, inplace=True))
block = nn.Sequential(*layers)
return block