| from contextlib import ExitStack
|
|
|
| import torch
|
| from torch import nn
|
| import torch.nn.functional as F
|
|
|
| from .basic_blocks import SeparableConv2d
|
| from .resnet import ResNetBackbone
|
| from ...model import ops
|
|
|
|
|
| class DeepLabV3Plus(nn.Module):
|
| def __init__(self, backbone='resnet50', norm_layer=nn.BatchNorm2d,
|
| backbone_norm_layer=None,
|
| ch=256,
|
| project_dropout=0.5,
|
| inference_mode=False,
|
| **kwargs):
|
| super(DeepLabV3Plus, self).__init__()
|
| if backbone_norm_layer is None:
|
| backbone_norm_layer = norm_layer
|
|
|
| self.backbone_name = backbone
|
| self.norm_layer = norm_layer
|
| self.backbone_norm_layer = backbone_norm_layer
|
| self.inference_mode = False
|
| self.ch = ch
|
| self.aspp_in_channels = 2048
|
| self.skip_project_in_channels = 256
|
|
|
| self._kwargs = kwargs
|
| if backbone == 'resnet34':
|
| self.aspp_in_channels = 512
|
| self.skip_project_in_channels = 64
|
|
|
| self.backbone = ResNetBackbone(backbone=self.backbone_name, pretrained_base=False,
|
| norm_layer=self.backbone_norm_layer, **kwargs)
|
|
|
| self.head = _DeepLabHead(in_channels=ch + 32, mid_channels=ch, out_channels=ch,
|
| norm_layer=self.norm_layer)
|
| self.skip_project = _SkipProject(self.skip_project_in_channels, 32, norm_layer=self.norm_layer)
|
| self.aspp = _ASPP(in_channels=self.aspp_in_channels,
|
| atrous_rates=[12, 24, 36],
|
| out_channels=ch,
|
| project_dropout=project_dropout,
|
| norm_layer=self.norm_layer)
|
|
|
| if inference_mode:
|
| self.set_prediction_mode()
|
|
|
| def load_pretrained_weights(self):
|
| pretrained = ResNetBackbone(backbone=self.backbone_name, pretrained_base=True,
|
| norm_layer=self.backbone_norm_layer, **self._kwargs)
|
| backbone_state_dict = self.backbone.state_dict()
|
| pretrained_state_dict = pretrained.state_dict()
|
|
|
| backbone_state_dict.update(pretrained_state_dict)
|
| self.backbone.load_state_dict(backbone_state_dict)
|
|
|
| if self.inference_mode:
|
| for param in self.backbone.parameters():
|
| param.requires_grad = False
|
|
|
| def set_prediction_mode(self):
|
| self.inference_mode = True
|
| self.eval()
|
|
|
| def forward(self, x):
|
| with ExitStack() as stack:
|
| if self.inference_mode:
|
| stack.enter_context(torch.no_grad())
|
|
|
| c1, _, c3, c4 = self.backbone(x)
|
| c1 = self.skip_project(c1)
|
|
|
| x = self.aspp(c4)
|
| x = F.interpolate(x, c1.size()[2:], mode='bilinear', align_corners=True)
|
| x = torch.cat((x, c1), dim=1)
|
| x = self.head(x)
|
|
|
| return x,
|
|
|
|
|
| class _SkipProject(nn.Module):
|
| def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d):
|
| super(_SkipProject, self).__init__()
|
| _activation = ops.select_activation_function("relu")
|
|
|
| self.skip_project = nn.Sequential(
|
| nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
|
| norm_layer(out_channels),
|
| _activation()
|
| )
|
|
|
| def forward(self, x):
|
| return self.skip_project(x)
|
|
|
|
|
| class _DeepLabHead(nn.Module):
|
| def __init__(self, out_channels, in_channels, mid_channels=256, norm_layer=nn.BatchNorm2d):
|
| super(_DeepLabHead, self).__init__()
|
|
|
| self.block = nn.Sequential(
|
| SeparableConv2d(in_channels=in_channels, out_channels=mid_channels, dw_kernel=3,
|
| dw_padding=1, activation='relu', norm_layer=norm_layer),
|
| SeparableConv2d(in_channels=mid_channels, out_channels=mid_channels, dw_kernel=3,
|
| dw_padding=1, activation='relu', norm_layer=norm_layer),
|
| nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=1)
|
| )
|
|
|
| def forward(self, x):
|
| return self.block(x)
|
|
|
|
|
| class _ASPP(nn.Module):
|
| def __init__(self, in_channels, atrous_rates, out_channels=256,
|
| project_dropout=0.5, norm_layer=nn.BatchNorm2d):
|
| super(_ASPP, self).__init__()
|
|
|
| b0 = nn.Sequential(
|
| nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=False),
|
| norm_layer(out_channels),
|
| nn.ReLU()
|
| )
|
|
|
| rate1, rate2, rate3 = tuple(atrous_rates)
|
| b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer)
|
| b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer)
|
| b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer)
|
| b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
|
|
|
| self.concurent = nn.ModuleList([b0, b1, b2, b3, b4])
|
|
|
| project = [
|
| nn.Conv2d(in_channels=5*out_channels, out_channels=out_channels,
|
| kernel_size=1, bias=False),
|
| norm_layer(out_channels),
|
| nn.ReLU()
|
| ]
|
| if project_dropout > 0:
|
| project.append(nn.Dropout(project_dropout))
|
| self.project = nn.Sequential(*project)
|
|
|
| def forward(self, x):
|
| x = torch.cat([block(x) for block in self.concurent], dim=1)
|
|
|
| return self.project(x)
|
|
|
|
|
| class _AsppPooling(nn.Module):
|
| def __init__(self, in_channels, out_channels, norm_layer):
|
| super(_AsppPooling, self).__init__()
|
|
|
| self.gap = nn.Sequential(
|
| nn.AdaptiveAvgPool2d((1, 1)),
|
| nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
|
| kernel_size=1, bias=False),
|
| norm_layer(out_channels),
|
| nn.ReLU()
|
| )
|
|
|
| def forward(self, x):
|
| pool = self.gap(x)
|
| return F.interpolate(pool, x.size()[2:], mode='bilinear', align_corners=True)
|
|
|
|
|
| def _ASPPConv(in_channels, out_channels, atrous_rate, norm_layer):
|
| block = nn.Sequential(
|
| nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
|
| kernel_size=3, padding=atrous_rate,
|
| dilation=atrous_rate, bias=False),
|
| norm_layer(out_channels),
|
| nn.ReLU()
|
| )
|
|
|
| return block
|
|
|
