| | |
| | import torch |
| | import torch.nn as nn |
| | from mmcv.cnn import build_activation_layer, build_conv_layer, build_norm_layer |
| | from mmengine.model import BaseModule |
| |
|
| | from mmseg.registry import MODELS |
| | from ..utils import resize |
| |
|
| |
|
| | class DownsamplerBlock(BaseModule): |
| | """Downsampler block of ERFNet. |
| | |
| | This module is a little different from basical ConvModule. |
| | The features from Conv and MaxPool layers are |
| | concatenated before BatchNorm. |
| | |
| | Args: |
| | in_channels (int): Number of input channels. |
| | out_channels (int): Number of output channels. |
| | conv_cfg (dict | None): Config of conv layers. |
| | Default: None. |
| | norm_cfg (dict | None): Config of norm layers. |
| | Default: dict(type='BN'). |
| | act_cfg (dict): Config of activation layers. |
| | Default: dict(type='ReLU'). |
| | init_cfg (dict or list[dict], optional): Initialization config dict. |
| | Default: None. |
| | """ |
| |
|
| | def __init__(self, |
| | in_channels, |
| | out_channels, |
| | conv_cfg=None, |
| | norm_cfg=dict(type='BN', eps=1e-3), |
| | act_cfg=dict(type='ReLU'), |
| | init_cfg=None): |
| | super().__init__(init_cfg=init_cfg) |
| | self.conv_cfg = conv_cfg |
| | self.norm_cfg = norm_cfg |
| | self.act_cfg = act_cfg |
| |
|
| | self.conv = build_conv_layer( |
| | self.conv_cfg, |
| | in_channels, |
| | out_channels - in_channels, |
| | kernel_size=3, |
| | stride=2, |
| | padding=1) |
| | self.pool = nn.MaxPool2d(kernel_size=2, stride=2) |
| | self.bn = build_norm_layer(self.norm_cfg, out_channels)[1] |
| | self.act = build_activation_layer(self.act_cfg) |
| |
|
| | def forward(self, input): |
| | conv_out = self.conv(input) |
| | pool_out = self.pool(input) |
| | pool_out = resize( |
| | input=pool_out, |
| | size=conv_out.size()[2:], |
| | mode='bilinear', |
| | align_corners=False) |
| | output = torch.cat([conv_out, pool_out], 1) |
| | output = self.bn(output) |
| | output = self.act(output) |
| | return output |
| |
|
| |
|
| | class NonBottleneck1d(BaseModule): |
| | """Non-bottleneck block of ERFNet. |
| | |
| | Args: |
| | channels (int): Number of channels in Non-bottleneck block. |
| | drop_rate (float): Probability of an element to be zeroed. |
| | Default 0. |
| | dilation (int): Dilation rate for last two conv layers. |
| | Default 1. |
| | num_conv_layer (int): Number of 3x1 and 1x3 convolution layers. |
| | Default 2. |
| | conv_cfg (dict | None): Config of conv layers. |
| | Default: None. |
| | norm_cfg (dict | None): Config of norm layers. |
| | Default: dict(type='BN'). |
| | act_cfg (dict): Config of activation layers. |
| | Default: dict(type='ReLU'). |
| | init_cfg (dict or list[dict], optional): Initialization config dict. |
| | Default: None. |
| | """ |
| |
|
| | def __init__(self, |
| | channels, |
| | drop_rate=0, |
| | dilation=1, |
| | num_conv_layer=2, |
| | conv_cfg=None, |
| | norm_cfg=dict(type='BN', eps=1e-3), |
| | act_cfg=dict(type='ReLU'), |
| | init_cfg=None): |
| | super().__init__(init_cfg=init_cfg) |
| |
|
| | self.conv_cfg = conv_cfg |
| | self.norm_cfg = norm_cfg |
| | self.act_cfg = act_cfg |
| | self.act = build_activation_layer(self.act_cfg) |
| |
|
| | self.convs_layers = nn.ModuleList() |
| | for conv_layer in range(num_conv_layer): |
| | first_conv_padding = (1, 0) if conv_layer == 0 else (dilation, 0) |
| | first_conv_dilation = 1 if conv_layer == 0 else (dilation, 1) |
| | second_conv_padding = (0, 1) if conv_layer == 0 else (0, dilation) |
| | second_conv_dilation = 1 if conv_layer == 0 else (1, dilation) |
| |
|
| | self.convs_layers.append( |
| | build_conv_layer( |
| | self.conv_cfg, |
| | channels, |
| | channels, |
| | kernel_size=(3, 1), |
| | stride=1, |
| | padding=first_conv_padding, |
| | bias=True, |
| | dilation=first_conv_dilation)) |
| | self.convs_layers.append(self.act) |
| | self.convs_layers.append( |
| | build_conv_layer( |
| | self.conv_cfg, |
| | channels, |
| | channels, |
| | kernel_size=(1, 3), |
| | stride=1, |
| | padding=second_conv_padding, |
| | bias=True, |
| | dilation=second_conv_dilation)) |
| | self.convs_layers.append( |
| | build_norm_layer(self.norm_cfg, channels)[1]) |
| | if conv_layer == 0: |
| | self.convs_layers.append(self.act) |
| | else: |
| | self.convs_layers.append(nn.Dropout(p=drop_rate)) |
| |
|
| | def forward(self, input): |
| | output = input |
| | for conv in self.convs_layers: |
| | output = conv(output) |
| | output = self.act(output + input) |
| | return output |
| |
|
| |
|
| | class UpsamplerBlock(BaseModule): |
| | """Upsampler block of ERFNet. |
| | |
| | Args: |
| | in_channels (int): Number of input channels. |
| | out_channels (int): Number of output channels. |
| | conv_cfg (dict | None): Config of conv layers. |
| | Default: None. |
| | norm_cfg (dict | None): Config of norm layers. |
| | Default: dict(type='BN'). |
| | act_cfg (dict): Config of activation layers. |
| | Default: dict(type='ReLU'). |
| | init_cfg (dict or list[dict], optional): Initialization config dict. |
| | Default: None. |
| | """ |
| |
|
| | def __init__(self, |
| | in_channels, |
| | out_channels, |
| | conv_cfg=None, |
| | norm_cfg=dict(type='BN', eps=1e-3), |
| | act_cfg=dict(type='ReLU'), |
| | init_cfg=None): |
| | super().__init__(init_cfg=init_cfg) |
| | self.conv_cfg = conv_cfg |
| | self.norm_cfg = norm_cfg |
| | self.act_cfg = act_cfg |
| |
|
| | self.conv = nn.ConvTranspose2d( |
| | in_channels=in_channels, |
| | out_channels=out_channels, |
| | kernel_size=3, |
| | stride=2, |
| | padding=1, |
| | output_padding=1, |
| | bias=True) |
| | self.bn = build_norm_layer(self.norm_cfg, out_channels)[1] |
| | self.act = build_activation_layer(self.act_cfg) |
| |
|
| | def forward(self, input): |
| | output = self.conv(input) |
| | output = self.bn(output) |
| | output = self.act(output) |
| | return output |
| |
|
| |
|
| | @MODELS.register_module() |
| | class ERFNet(BaseModule): |
| | """ERFNet backbone. |
| | |
| | This backbone is the implementation of `ERFNet: Efficient Residual |
| | Factorized ConvNet for Real-time SemanticSegmentation |
| | <https://ieeexplore.ieee.org/document/8063438>`_. |
| | |
| | Args: |
| | in_channels (int): The number of channels of input |
| | image. Default: 3. |
| | enc_downsample_channels (Tuple[int]): Size of channel |
| | numbers of various Downsampler block in encoder. |
| | Default: (16, 64, 128). |
| | enc_stage_non_bottlenecks (Tuple[int]): Number of stages of |
| | Non-bottleneck block in encoder. |
| | Default: (5, 8). |
| | enc_non_bottleneck_dilations (Tuple[int]): Dilation rate of each |
| | stage of Non-bottleneck block of encoder. |
| | Default: (2, 4, 8, 16). |
| | enc_non_bottleneck_channels (Tuple[int]): Size of channel |
| | numbers of various Non-bottleneck block in encoder. |
| | Default: (64, 128). |
| | dec_upsample_channels (Tuple[int]): Size of channel numbers of |
| | various Deconvolution block in decoder. |
| | Default: (64, 16). |
| | dec_stages_non_bottleneck (Tuple[int]): Number of stages of |
| | Non-bottleneck block in decoder. |
| | Default: (2, 2). |
| | dec_non_bottleneck_channels (Tuple[int]): Size of channel |
| | numbers of various Non-bottleneck block in decoder. |
| | Default: (64, 16). |
| | drop_rate (float): Probability of an element to be zeroed. |
| | Default 0.1. |
| | """ |
| |
|
| | def __init__(self, |
| | in_channels=3, |
| | enc_downsample_channels=(16, 64, 128), |
| | enc_stage_non_bottlenecks=(5, 8), |
| | enc_non_bottleneck_dilations=(2, 4, 8, 16), |
| | enc_non_bottleneck_channels=(64, 128), |
| | dec_upsample_channels=(64, 16), |
| | dec_stages_non_bottleneck=(2, 2), |
| | dec_non_bottleneck_channels=(64, 16), |
| | dropout_ratio=0.1, |
| | conv_cfg=None, |
| | norm_cfg=dict(type='BN', requires_grad=True), |
| | act_cfg=dict(type='ReLU'), |
| | init_cfg=None): |
| |
|
| | super().__init__(init_cfg=init_cfg) |
| | assert len(enc_downsample_channels) \ |
| | == len(dec_upsample_channels)+1, 'Number of downsample\ |
| | block of encoder does not \ |
| | match number of upsample block of decoder!' |
| | assert len(enc_downsample_channels) \ |
| | == len(enc_stage_non_bottlenecks)+1, 'Number of \ |
| | downsample block of encoder does not match \ |
| | number of Non-bottleneck block of encoder!' |
| | assert len(enc_downsample_channels) \ |
| | == len(enc_non_bottleneck_channels)+1, 'Number of \ |
| | downsample block of encoder does not match \ |
| | number of channels of Non-bottleneck block of encoder!' |
| | assert enc_stage_non_bottlenecks[-1] \ |
| | % len(enc_non_bottleneck_dilations) == 0, 'Number of \ |
| | Non-bottleneck block of encoder does not match \ |
| | number of Non-bottleneck block of encoder!' |
| | assert len(dec_upsample_channels) \ |
| | == len(dec_stages_non_bottleneck), 'Number of \ |
| | upsample block of decoder does not match \ |
| | number of Non-bottleneck block of decoder!' |
| | assert len(dec_stages_non_bottleneck) \ |
| | == len(dec_non_bottleneck_channels), 'Number of \ |
| | Non-bottleneck block of decoder does not match \ |
| | number of channels of Non-bottleneck block of decoder!' |
| |
|
| | self.in_channels = in_channels |
| | self.enc_downsample_channels = enc_downsample_channels |
| | self.enc_stage_non_bottlenecks = enc_stage_non_bottlenecks |
| | self.enc_non_bottleneck_dilations = enc_non_bottleneck_dilations |
| | self.enc_non_bottleneck_channels = enc_non_bottleneck_channels |
| | self.dec_upsample_channels = dec_upsample_channels |
| | self.dec_stages_non_bottleneck = dec_stages_non_bottleneck |
| | self.dec_non_bottleneck_channels = dec_non_bottleneck_channels |
| | self.dropout_ratio = dropout_ratio |
| |
|
| | self.encoder = nn.ModuleList() |
| | self.decoder = nn.ModuleList() |
| |
|
| | self.conv_cfg = conv_cfg |
| | self.norm_cfg = norm_cfg |
| | self.act_cfg = act_cfg |
| |
|
| | self.encoder.append( |
| | DownsamplerBlock(self.in_channels, enc_downsample_channels[0])) |
| |
|
| | for i in range(len(enc_downsample_channels) - 1): |
| | self.encoder.append( |
| | DownsamplerBlock(enc_downsample_channels[i], |
| | enc_downsample_channels[i + 1])) |
| | |
| | if i == len(enc_downsample_channels) - 2: |
| | iteration_times = int(enc_stage_non_bottlenecks[-1] / |
| | len(enc_non_bottleneck_dilations)) |
| | for j in range(iteration_times): |
| | for k in range(len(enc_non_bottleneck_dilations)): |
| | self.encoder.append( |
| | NonBottleneck1d(enc_downsample_channels[-1], |
| | self.dropout_ratio, |
| | enc_non_bottleneck_dilations[k])) |
| | else: |
| | for j in range(enc_stage_non_bottlenecks[i]): |
| | self.encoder.append( |
| | NonBottleneck1d(enc_downsample_channels[i + 1], |
| | self.dropout_ratio)) |
| |
|
| | for i in range(len(dec_upsample_channels)): |
| | if i == 0: |
| | self.decoder.append( |
| | UpsamplerBlock(enc_downsample_channels[-1], |
| | dec_non_bottleneck_channels[i])) |
| | else: |
| | self.decoder.append( |
| | UpsamplerBlock(dec_non_bottleneck_channels[i - 1], |
| | dec_non_bottleneck_channels[i])) |
| | for j in range(dec_stages_non_bottleneck[i]): |
| | self.decoder.append( |
| | NonBottleneck1d(dec_non_bottleneck_channels[i])) |
| |
|
| | def forward(self, x): |
| | for enc in self.encoder: |
| | x = enc(x) |
| | for dec in self.decoder: |
| | x = dec(x) |
| | return [x] |
| |
|
