| | import torch |
| | from torch import nn |
| | from torch.nn import functional as F |
| |
|
| | from .utils import _SimpleSegmentationModel |
| |
|
| |
|
| | __all__ = ["DeepLabV3"] |
| |
|
| |
|
| | class DeepLabV3(_SimpleSegmentationModel): |
| | """ |
| | Implements DeepLabV3 model from |
| | `"Rethinking Atrous Convolution for Semantic Image Segmentation" |
| | <https://arxiv.org/abs/1706.05587>`_. |
| | |
| | Arguments: |
| | backbone (nn.Module): the network used to compute the features for the model. |
| | The backbone should return an OrderedDict[Tensor], with the key being |
| | "out" for the last feature map used, and "aux" if an auxiliary classifier |
| | is used. |
| | classifier (nn.Module): module that takes the "out" element returned from |
| | the backbone and returns a dense prediction. |
| | aux_classifier (nn.Module, optional): auxiliary classifier used during training |
| | """ |
| | pass |
| |
|
| | class DeepLabHeadV3Plus(nn.Module): |
| | def __init__(self, in_channels, low_level_channels, num_classes, aspp_dilate=[12, 24, 36]): |
| | super(DeepLabHeadV3Plus, self).__init__() |
| | self.project = nn.Sequential( |
| | nn.Conv2d(low_level_channels, 48, 1, bias=False), |
| | nn.BatchNorm2d(48), |
| | nn.ReLU(inplace=True), |
| | ) |
| |
|
| | self.aspp = ASPP(in_channels, aspp_dilate) |
| |
|
| | self.classifier = nn.Sequential( |
| | nn.Conv2d(304, 256, 3, padding=1, bias=False), |
| | nn.BatchNorm2d(256), |
| | nn.ReLU(inplace=True), |
| | nn.Conv2d(256, num_classes, 1) |
| | ) |
| | self._init_weight() |
| |
|
| | def forward(self, feature): |
| | low_level_feature = self.project( feature['low_level'] ) |
| | output_feature = self.aspp(feature['out']) |
| | output_feature = F.interpolate(output_feature, size=low_level_feature.shape[2:], mode='bilinear', align_corners=False) |
| | return self.classifier( torch.cat( [ low_level_feature, output_feature ], dim=1 ) ) |
| | |
| | def _init_weight(self): |
| | for m in self.modules(): |
| | if isinstance(m, nn.Conv2d): |
| | nn.init.kaiming_normal_(m.weight) |
| | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): |
| | nn.init.constant_(m.weight, 1) |
| | nn.init.constant_(m.bias, 0) |
| |
|
| | class DeepLabHead(nn.Module): |
| | def __init__(self, in_channels, num_classes, aspp_dilate=[12, 24, 36]): |
| | super(DeepLabHead, self).__init__() |
| |
|
| | self.classifier = nn.Sequential( |
| | ASPP(in_channels, aspp_dilate), |
| | nn.Conv2d(256, 256, 3, padding=1, bias=False), |
| | nn.BatchNorm2d(256), |
| | nn.ReLU(inplace=True), |
| | nn.Conv2d(256, num_classes, 1) |
| | ) |
| | self._init_weight() |
| |
|
| | def forward(self, feature): |
| | return self.classifier( feature['out'] ) |
| |
|
| | def _init_weight(self): |
| | for m in self.modules(): |
| | if isinstance(m, nn.Conv2d): |
| | nn.init.kaiming_normal_(m.weight) |
| | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): |
| | nn.init.constant_(m.weight, 1) |
| | nn.init.constant_(m.bias, 0) |
| |
|
| | class AtrousSeparableConvolution(nn.Module): |
| | """ Atrous Separable Convolution |
| | """ |
| | def __init__(self, in_channels, out_channels, kernel_size, |
| | stride=1, padding=0, dilation=1, bias=True): |
| | super(AtrousSeparableConvolution, self).__init__() |
| | self.body = nn.Sequential( |
| | |
| | nn.Conv2d( in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, groups=in_channels ), |
| | |
| | nn.Conv2d( in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), |
| | ) |
| | |
| | self._init_weight() |
| |
|
| | def forward(self, x): |
| | return self.body(x) |
| |
|
| | def _init_weight(self): |
| | for m in self.modules(): |
| | if isinstance(m, nn.Conv2d): |
| | nn.init.kaiming_normal_(m.weight) |
| | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): |
| | nn.init.constant_(m.weight, 1) |
| | nn.init.constant_(m.bias, 0) |
| |
|
| | class ASPPConv(nn.Sequential): |
| | def __init__(self, in_channels, out_channels, dilation): |
| | modules = [ |
| | nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False), |
| | nn.BatchNorm2d(out_channels), |
| | nn.ReLU(inplace=True) |
| | ] |
| | super(ASPPConv, self).__init__(*modules) |
| |
|
| | class ASPPPooling(nn.Sequential): |
| | def __init__(self, in_channels, out_channels): |
| | super(ASPPPooling, self).__init__( |
| | nn.AdaptiveAvgPool2d(1), |
| | nn.Conv2d(in_channels, out_channels, 1, bias=False), |
| | nn.BatchNorm2d(out_channels), |
| | nn.ReLU(inplace=True)) |
| |
|
| | def forward(self, x): |
| | size = x.shape[-2:] |
| | x = super(ASPPPooling, self).forward(x) |
| | return F.interpolate(x, size=size, mode='bilinear', align_corners=False) |
| |
|
| | class ASPP(nn.Module): |
| | def __init__(self, in_channels, atrous_rates): |
| | super(ASPP, self).__init__() |
| | out_channels = 256 |
| | modules = [] |
| | modules.append(nn.Sequential( |
| | nn.Conv2d(in_channels, out_channels, 1, bias=False), |
| | nn.BatchNorm2d(out_channels), |
| | nn.ReLU(inplace=True))) |
| |
|
| | rate1, rate2, rate3 = tuple(atrous_rates) |
| | modules.append(ASPPConv(in_channels, out_channels, rate1)) |
| | modules.append(ASPPConv(in_channels, out_channels, rate2)) |
| | modules.append(ASPPConv(in_channels, out_channels, rate3)) |
| | modules.append(ASPPPooling(in_channels, out_channels)) |
| |
|
| | self.convs = nn.ModuleList(modules) |
| |
|
| | self.project = nn.Sequential( |
| | nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), |
| | nn.BatchNorm2d(out_channels), |
| | nn.ReLU(inplace=True), |
| | nn.Dropout(0.1),) |
| |
|
| | def forward(self, x): |
| | res = [] |
| | for conv in self.convs: |
| | res.append(conv(x)) |
| | res = torch.cat(res, dim=1) |
| | return self.project(res) |
| |
|
| |
|
| |
|
| | def convert_to_separable_conv(module): |
| | new_module = module |
| | if isinstance(module, nn.Conv2d) and module.kernel_size[0]>1: |
| | new_module = AtrousSeparableConvolution(module.in_channels, |
| | module.out_channels, |
| | module.kernel_size, |
| | module.stride, |
| | module.padding, |
| | module.dilation, |
| | module.bias) |
| | for name, child in module.named_children(): |
| | new_module.add_module(name, convert_to_separable_conv(child)) |
| | return new_module |
