pytorch构建deeplabv3+

本文主要是介绍pytorch构建deeplabv3+，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

DeepLab v3+ 是DeepLab语义分割系列网络的最新作，其前作有 DeepLab v1，v2, v3, 在最新作中，Liang-Chieh Chen等人通过encoder-decoder进行多尺度信息的融合，同时保留了原来的空洞卷积和ASSP层，其骨干网络使用了Xception模型，提高了语义分割的健壮性和运行速率。其在Pascal VOC上达到了 89.0% 的mIoU，在Cityscape上也取得了 82.1%的好成绩，下图展示了DeepLab v3+的基本结构：

请添加图片描述
其实在DCNN中主要是做一个特征提取，至于采用哪个网络做backbone具体问题具体对待，在这里我才用的是mobilenetv2(只是将deepwise_conv中添加了dilation，添加空洞卷积是为了增大感受野)

网络结构分为Encode部分和decoder部分
先看encoder部分：
请添加图片描述
接在DCNN后面的实际上就是一个ASPP结构（采用不同的采样率来对特征图做空洞卷积），然后再将对应的结果进行拼接，需要注意的是传入ASPP结构的是DCNN得到的高层特征图image Pooling部分其实会改变特征图的尺寸，所以可以通过使用双线插值（为什么采用双线插值，因为简单）或者其他方式保证经过ASPP结构的各个特征图尺寸相同，最后再进行拼接
请添加图片描述

再看decoder部分请添加图片描述
decoder部分首先会对传入的低层特征图进行通道调整，然后与encoder传入的特征图进行拼接，注意encoder传入的特征图需要经过上采样处理（维持与低层特征图相同的尺寸），最后输出部分只需要将尺寸还原到输入图片的尺寸就行了

import torch
import torch.nn as nn
import torch.functional as Fclass ASPP(nn.Module):def __init__(self, feature, atrous):super(ASPP, self).__init__()self.feature = featureself.Conv1 = _Deepwise_Conv(in_channels=feature.size()[1], out_channels=256, use_bias=False)self.Conv_rate1 = _Deepwise_Conv(in_channels=feature.size()[1], out_channels=256, rate=atrous[0],padding=atrous[0], use_bias=False)self.Conv_rate2 = _Deepwise_Conv(in_channels=feature.size()[1], out_channels=256, rate=atrous[1],padding=atrous[1], use_bias=False)self.Conv_rate3 = _Deepwise_Conv(in_channels=feature.size()[1], out_channels=256, rate=atrous[2],padding=atrous[2], use_bias=False)self.globalAvgPoolAndConv = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),Conv(in_channels=320, out_channels=256, kernel_size=1, stride=1, use_bias=False),)self.Conv4 = Conv(in_channels=256 * 5, out_channels=256, kernel_size=1, stride=1, use_bias=False)self.dropout = nn.Dropout(p=0.1)def forward(self):f1 = self.Conv1(self.feature.clone())f2 = self.Conv_rate1(self.feature.clone())f3 = self.Conv_rate2(self.feature.clone())f4 = self.Conv_rate3(self.feature.clone())f5 = self.globalAvgPoolAndConv(self.feature.clone())f5 = F.interpolate(f5, size=(self.feature.size(2), self.feature.size(3)), mode='bilinear')x = torch.cat([f1, f2, f3, f4, f5], dim=1)x = self.Conv4(x)x = self.dropout(x)class Deeplabv3(nn.Module):def __init__(self, feature, atrous, skip1, num_class):super(Deeplabv3, self).__init__()self.num_class = num_classself.feature = ASPP(atrous=atrous, feature=feature).forward()self.skip1 = skip1self.encoder = ASPP(atrous=atrous, feature=feature)self.Conv1 = Conv(in_channels=skip1.size()[1], out_channels=48, kernel_size=1,strip=1, use_bias=False)self.Conv2 = _Deepwise_Conv(in_channels=48 + 256, out_channels=256, use_bias=False)self.ConvNUM = Conv(in_channels=256, out_channels=num_class, kernel_size=1, use_bias=False)def forward(self, input_img):skip1 = self.Conv1(self.skip1)feature = F.interpolate(self.feature, size=(skip1.size()[2], skip1.size()[3]), mode='bilinear')skip1 = torch.cat([skip1, feature], dim=1)skip1 = self.Conv2(skip1)skip1 = self.ConvNUM(skip1)skip1 = F.interpolate(skip1, size=(input_img.size()[2], input_img.size()[3]))return F.softmax(skip1,dim=1)class _bottlenet(nn.Module):def __init__(self, in_channels, out_channels, rate=1, expand_ratio=1, stride=1):super(_bottlenet, self).__init__()# 步长为2以及前后通道数不同就不进行残差堆叠self.use_res_connect = (stride == 1) and (in_channels == out_channels)self.features = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=in_channels * expand_ratio, kernel_size=1),nn.BatchNorm2d(num_features=in_channels * expand_ratio),nn.ReLU6(inplace=True),nn.Conv2d(in_channels=in_channels * expand_ratio, out_channels=in_channels * expand_ratio, kernel_size=3, stride=stride,padding=rate, dilation=(rate, rate)),nn.BatchNorm2d(num_features=in_channels * expand_ratio),nn.ReLU6(inplace=True),nn.Conv2d(in_channels=in_channels * expand_ratio, out_channels=out_channels, stride=1, kernel_size=1,padding=0),nn.BatchNorm2d(num_features=out_channels),nn.ReLU6(inplace=True),)# self.change = nn.Conv2d()def forward(self, x):x_clone = x.clone()x = self.features(x)#         print(x.size())if self.use_res_connect:#             print("="*10)#             print(x.size())#             print(x_clone.size())x.add_(x_clone)return xclass get_mobilenetv2_encoder(nn.Module):def __init__(self, downsamp_factor=8, num_classes=3):super(get_mobilenetv2_encoder, self).__init__()if downsamp_factor == 8:self.atrous_rates = (12, 24, 36)block4_dilation = 2block5_dilation = 4block4_stride = 1else:self.atrous_rates = (6, 12, 18)block4_dilation = 1block5_dilation = 2block4_stride = 2self.features = []self.features.append(nn.Conv2d(in_channels=3, out_channels=32, kernel_size=(3, 3), padding=1, stride=2))self.features.append(nn.BatchNorm2d(num_features=32))self.features.append(nn.ReLU6(inplace=True))# ------  3 ------# block1self.features.append(_bottlenet(in_channels=32, out_channels=16, expand_ratio=1, stride=1))# block2# [t, c, n, s] = [6, 24, 2, 2]self.features.append(_bottlenet(in_channels=16, out_channels=24, expand_ratio=6, stride=2))self.features.append(_bottlenet(in_channels=24, out_channels=24, expand_ratio=6, stride=1))# ------  6  -----# block3# [t, c, n, s] = [6, 32, 3, 2]self.features.append(_bottlenet(in_channels=24, out_channels=32, expand_ratio=6, stride=2))for i in range(2):self.features.append(_bottlenet(in_channels=32, out_channels=32, expand_ratio=6))# ------  9  ------# block4# [t, c, n, s] = [6, 64, 4, 2]self.features.append(_bottlenet(in_channels=32, out_channels=64, expand_ratio=6, stride=block4_stride))for i in range(3):self.features.append(_bottlenet(in_channels=64, out_channels=64, expand_ratio=6, rate=block4_dilation))# ------  13  ------# block5# [t, c, n, s] = [6, 96, 3, 1]self.features.append(_bottlenet(in_channels=64, out_channels=96, expand_ratio=6, rate=block4_dilation))for i in range(2):self.features.append(_bottlenet(in_channels=96, out_channels=96, expand_ratio=6, rate=block4_dilation))# [t, c, n, s] = [6, 160, 3, 2]# block6self.features.append(_bottlenet(in_channels=96, out_channels=160, expand_ratio=6, stride=1))for i in range(2):self.features.append(_bottlenet(in_channels=160, out_channels=160, expand_ratio=6))# [t, c, n, s] = [6, 160, 3, 2]self.features.append(_bottlenet(in_channels=160, out_channels=320, expand_ratio=6))self.features = nn.Sequential(*self.features)def forward(self, x):skip1 = Nonefor i, op in enumerate(self.features, 0):x = op(x)if i == 5:skip1 = x.clone()return x, self.atrous_rates, skip1class pool_block(nn.Module):def __init__(self, f, stride):super(pool_block, self).__init__()in_channels = f.size()[1]kernel_size = strideself.features = nn.Sequential(nn.AvgPool2d(kernel_size=kernel_size, stride=kernel_size, padding=kernel_size // 2),nn.Conv2d(in_channels=in_channels, out_channels=512, kernel_size=1, stride=1, bias=False),nn.BatchNorm2d(num_features=512),nn.ReLU6(inplace=True),nn.Upsample(size=(INPUT_SIZE, INPUT_SIZE), mode="bilinear"))def forward(self, x):x = self.features(x)return xclass _Deepwise_Conv(nn.Module):def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, rate=1, use_bias=False):super(_Deepwise_Conv, self).__init__()self.conv1 = Conv(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,stride=stride, padding=padding, dilation=rate, use_bias=use_bias)self.conv2 = Conv(in_channels=in_channels, out_channels=out_channels, kernel_size=1,stride=1, padding=0, use_bias=use_bias)def forward(self, x):return self.conv2(self.conv1(x))class Conv(nn.Module):'''nn.Conv2d + Batchnormlizetion + ReLU6'''def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, use_bias=False):super(Conv, self).__init__()self.features = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,stride=stride, padding=padding, dilation=dilation, bias=use_bias),nn.BatchNorm2d(num_features=out_channels),nn.ReLU6(),)def forward(self, x):return self.features(x)