| 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 |
