| from collections import OrderedDict
|
|
|
| import torch.nn as nn
|
|
|
| from .bn import ABN, ACT_LEAKY_RELU, ACT_ELU, ACT_NONE
|
| import torch.nn.functional as functional
|
|
|
|
|
| class ResidualBlock(nn.Module):
|
| """Configurable residual block
|
|
|
| Parameters
|
| ----------
|
| in_channels : int
|
| Number of input channels.
|
| channels : list of int
|
| Number of channels in the internal feature maps. Can either have two or three elements: if three construct
|
| a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then
|
| `3 x 3` then `1 x 1` convolutions.
|
| stride : int
|
| Stride of the first `3 x 3` convolution
|
| dilation : int
|
| Dilation to apply to the `3 x 3` convolutions.
|
| groups : int
|
| Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with
|
| bottleneck blocks.
|
| norm_act : callable
|
| Function to create normalization / activation Module.
|
| dropout: callable
|
| Function to create Dropout Module.
|
| """
|
|
|
| def __init__(self,
|
| in_channels,
|
| channels,
|
| stride=1,
|
| dilation=1,
|
| groups=1,
|
| norm_act=ABN,
|
| dropout=None):
|
| super(ResidualBlock, self).__init__()
|
|
|
|
|
| if len(channels) != 2 and len(channels) != 3:
|
| raise ValueError("channels must contain either two or three values")
|
| if len(channels) == 2 and groups != 1:
|
| raise ValueError("groups > 1 are only valid if len(channels) == 3")
|
|
|
| is_bottleneck = len(channels) == 3
|
| need_proj_conv = stride != 1 or in_channels != channels[-1]
|
|
|
| if not is_bottleneck:
|
| bn2 = norm_act(channels[1])
|
| bn2.activation = ACT_NONE
|
| layers = [
|
| ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False,
|
| dilation=dilation)),
|
| ("bn1", norm_act(channels[0])),
|
| ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
|
| dilation=dilation)),
|
| ("bn2", bn2)
|
| ]
|
| if dropout is not None:
|
| layers = layers[0:2] + [("dropout", dropout())] + layers[2:]
|
| else:
|
| bn3 = norm_act(channels[2])
|
| bn3.activation = ACT_NONE
|
| layers = [
|
| ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=1, padding=0, bias=False)),
|
| ("bn1", norm_act(channels[0])),
|
| ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=stride, padding=dilation, bias=False,
|
| groups=groups, dilation=dilation)),
|
| ("bn2", norm_act(channels[1])),
|
| ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False)),
|
| ("bn3", bn3)
|
| ]
|
| if dropout is not None:
|
| layers = layers[0:4] + [("dropout", dropout())] + layers[4:]
|
| self.convs = nn.Sequential(OrderedDict(layers))
|
|
|
| if need_proj_conv:
|
| self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False)
|
| self.proj_bn = norm_act(channels[-1])
|
| self.proj_bn.activation = ACT_NONE
|
|
|
| def forward(self, x):
|
| if hasattr(self, "proj_conv"):
|
| residual = self.proj_conv(x)
|
| residual = self.proj_bn(residual)
|
| else:
|
| residual = x
|
| x = self.convs(x) + residual
|
|
|
| if self.convs.bn1.activation == ACT_LEAKY_RELU:
|
| return functional.leaky_relu(x, negative_slope=self.convs.bn1.slope, inplace=True)
|
| elif self.convs.bn1.activation == ACT_ELU:
|
| return functional.elu(x, inplace=True)
|
| else:
|
| return x
|
|
|
|
|
| class IdentityResidualBlock(nn.Module):
|
| def __init__(self,
|
| in_channels,
|
| channels,
|
| stride=1,
|
| dilation=1,
|
| groups=1,
|
| norm_act=ABN,
|
| dropout=None):
|
| """Configurable identity-mapping residual block
|
|
|
| Parameters
|
| ----------
|
| in_channels : int
|
| Number of input channels.
|
| channels : list of int
|
| Number of channels in the internal feature maps. Can either have two or three elements: if three construct
|
| a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then
|
| `3 x 3` then `1 x 1` convolutions.
|
| stride : int
|
| Stride of the first `3 x 3` convolution
|
| dilation : int
|
| Dilation to apply to the `3 x 3` convolutions.
|
| groups : int
|
| Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with
|
| bottleneck blocks.
|
| norm_act : callable
|
| Function to create normalization / activation Module.
|
| dropout: callable
|
| Function to create Dropout Module.
|
| """
|
| super(IdentityResidualBlock, self).__init__()
|
|
|
|
|
| if len(channels) != 2 and len(channels) != 3:
|
| raise ValueError("channels must contain either two or three values")
|
| if len(channels) == 2 and groups != 1:
|
| raise ValueError("groups > 1 are only valid if len(channels) == 3")
|
|
|
| is_bottleneck = len(channels) == 3
|
| need_proj_conv = stride != 1 or in_channels != channels[-1]
|
|
|
| self.bn1 = norm_act(in_channels)
|
| if not is_bottleneck:
|
| layers = [
|
| ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False,
|
| dilation=dilation)),
|
| ("bn2", norm_act(channels[0])),
|
| ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
|
| dilation=dilation))
|
| ]
|
| if dropout is not None:
|
| layers = layers[0:2] + [("dropout", dropout())] + layers[2:]
|
| else:
|
| layers = [
|
| ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=stride, padding=0, bias=False)),
|
| ("bn2", norm_act(channels[0])),
|
| ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
|
| groups=groups, dilation=dilation)),
|
| ("bn3", norm_act(channels[1])),
|
| ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False))
|
| ]
|
| if dropout is not None:
|
| layers = layers[0:4] + [("dropout", dropout())] + layers[4:]
|
| self.convs = nn.Sequential(OrderedDict(layers))
|
|
|
| if need_proj_conv:
|
| self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False)
|
|
|
| def forward(self, x):
|
| if hasattr(self, "proj_conv"):
|
| bn1 = self.bn1(x)
|
| shortcut = self.proj_conv(bn1)
|
| else:
|
| shortcut = x.clone()
|
| bn1 = self.bn1(x)
|
|
|
| out = self.convs(bn1)
|
| out.add_(shortcut)
|
|
|
| return out
|
|
|
