gcan pytorch实现_DeepLabv2: 论文阅读与Pytorch实现

作 者： 心有宝宝人自圆

声 明： 欢迎转载本文中的图片或文字，请说明出处

写在前面

自从FCN提出以来，越来越多的语义分割任务开始采用采用全卷积网络结构，随着FCN结构使用的增加，研究人员发先了其结构天生的缺陷极大的限制了分割的准确度：CNNs在high-level (large scale) tasks中取得了十分优异的成绩，这得益于局部空间不变性(主要是池化层增大了感受野，也丢弃了部分细节信息)使得网络能够学习到层次化的抽象信息，但这却恰恰不利于low-level (small scale) tasks

DeepLabv2的作者延续DeepLabv1的atrous算法和denseCRF的后处理边缘优化，并在Spatial Pyramid Pooling结合多尺度特征图以最大化信息保留和计算节约的特性启发之下，设计了Atrous Spatial Pyramid Pooling(ASPP)来对抗这种细节丢失的问题，并总结为如下三个主要贡献：

带洞的卷积，atrous算法

ASPP能够捕捉多尺度目标的特征

Fully connected CRF的后处理过程

1. Introduction

DCNNs在high-level vision tasks(如图像分类、目标检测等)取得优异得表现，这些工作都有共同的主题：end-to-end训练的方法比人工的特征工程方法更优。这得益于CNNs内在的局部空间不变性，然而对应low-level vision tasks(语义分割)来说，我们需要准确的位置信息而非空间信息抽象后的层次化信息。

DCNNs应用于low-level vision tasks主要难点是：

信号的下采样使细节信息丢失

网络对多尺度目标的泛化能力较差

空间的局部不变性导致位置信息不再准确

下采样问题是池化层和stride的联合影响造成，是准确性和速度、空间消耗的权衡(即过大的卷积核运算极度耗时、参数过大)。其目的是为了使较小的卷积核能够去学习空间中有用的信息(因此需要增大感受野)，但这种下采样必然造成信息的损失。为了在不造成信息损失的情况下增大感受野，作者使用了带洞的卷积(下均称atrous方法)，由此来显式地控制感受野的大小，atrous卷积设计的核心理念是临近的像素携带的信息相似。

DILATED CONVOLUTIONS with kernel size 3x3, dilation=2

真实的检测目标往往是多尺度的，而我们的训练集一般仅提供了较大尺度的目标。为了解决这一问题，主流的做法是将同一输入图片按不同尺度缩小输入相同参数的网络(对于缩小的图片，同样的kernel_size能获得更大的感受野，间接地达到了小尺度目标地效果)，再将特征图或得分图结合起来得到最终得得分。这样的操作的确带来了改进，当使训练更加地耗时。作者受Spatial Pyramid Pooling地启发，设计出了对特定特征图进行多尺度特征上采样提取的高效计算方法：ASPP。

1.PNG

局部空间不变性是classifier获得以对象为中心的决策的要求，主要还是由于池化层得作用只保留了局部空间中最重要的信息，作者使用Fully connected CRF(后称DenseCRF)进行全卷积网络训练完成后的后处理，DenseCRF能够在满足长程依赖性的同时捕获细节边缘信息。

2. METHODS

文章大部分内容和DeepLabv1十分相似，最主要的不同就是提出了ASPP。若对Atrous算法和DenseCRF算法有所需求，可以参考DeepLabv1，这里不再赘述

2.1 Multiscale Image Representations

DCNN具有从训练集中提取复杂的抽象特征的强大能力，而这写特征存在隐式的尺度，因此显式地考虑输入尺度能够提升网络模型对于各种尺度特征的泛化能力。为了解决语义分割中的多尺度目标难题，作者测试了两种处理方法：(1)将输入图片进行多次降尺度传入相同参数的网络，并将最终的得分图经过双线性插值后结合起来；(2)ASPP。

我们先来介绍一下Spatial Pyramid Pooling(SPP)

SPP将最后的MaxPool更改成了多尺度的AdaptiveMaxPool(为了组成固定长度向量传给fc层)，将conv5的特征图(Vgg-16)分别传入AdaptiveMaxPool(4), AdaptiveMaxPool(2), AdaptiveMaxPool(1)(仅是描述上图)，得到了不同尺度的特征图。实验结果表明能够很好的增强网络对于多尺度目标的鲁棒性、提升速度、接收任意大小的输入，此外由于AdaptiveMaxPool(1)的存在使得模型能减少过拟合。SPP虽然获得了更大的感受野，但丢失了空间信息，与atrous应用的场景类似，作者将池化层更改为了atrous卷积层，在获得更大感受野的同时保留了原有的空间信息；卷积层同样拥有fc层的特性，ASPP就能代替SPP+fc。

ASPP

作者使用拥有不同采样率(感受野)的atrous卷积层，使特征图并行通过各个atrous卷积层提取出各个尺度的中心像素点特征图，最后将各个尺度的特征图结合起来得到最后得得分图，能增强网络对多尺度目标的鲁棒性，提升网络的性能。

2.2 训练方案改进

2.2.1 Learning rate policy

之前作者适用的方法是“step” learning rate policy(达到2000iter的时候使，learning rate衰减为原来的0.1)；在文章中作者表明“poly” learning rate policy(

)能够得到更好的表现。

2.2.2 ASPP

ASPP是LargeFOV(具体可参考DeepLabv1的网络设计部分)的泛化，作者提供了ASPP_S( atrous rate = {2, 4, 8, 12})和ASPP_L( atrous rate = {6, 12, 18, 24})两种感受野，ASPP_L在参数更少的优势下，能够得到于ASPP_S相类似的效果

2.2.3 Deeper Networks and Multiscale Processing

作者使用预训练的ResNet-101代替Vgg-16，并采用2.1所述的将输入图片进行多次降尺度传入相同参数的网络的方法来获得更好的表现。

3. My Code

3.1 ResNet-101 Base

import torch.nn as nn

import torch.nn.functional as F

import torch

import torchvision

torchvision.models.segmentation.deeplabv3_resnet101()

BOTTLENECK_EXPANSION = 4

class ConvBnReLUBlock(nn.Sequential):

def __init__(self, index, in_channels, out_channels, kernel_size, stride, padding, dilation, relu=True):

super(ConvBnReLUBlock, self).__init__()

# 注意conv层中均无bias

self.add_module('conv' + str(index),

nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=False))

self.add_module('bn' + str(index), nn.BatchNorm2d(out_channels, eps=1e-5, momentum=0.1))

if relu:

self.add_module("relu" + str(index), nn.ReLU(inplace=True))

class Bottlenenck(nn.Module):

def __init__(self, in_channels, out_channels, stride, dilation, downsample):

super(Bottlenenck, self).__init__()

mid_channels = out_channels // BOTTLENECK_EXPANSION

# 利用stride进行下采样 -> f(x) - x

self.reduce = ConvBnReLUBlock(1, in_channels, mid_channels, 1, stride, 0, 1, True)

# padding = dilation * (kernel_size - 1)//2 ,保持形状不变

self.conv3x3 = ConvBnReLUBlock(2, mid_channels, mid_channels, 3, 1, dilation, dilation, True)

# 增加通道数但形状不变

self.increase = ConvBnReLUBlock(3, mid_channels, out_channels, 1, 1, 0, 1, False)

# 当in_channel != out_channel时先下采样 x

self.shortcut = nn.Sequential(ConvBnReLUBlock(4, in_channels, out_channels, 1, stride, 0, 1,

False) if downsample else nn.Identity())

def forward(self, x):

out = self.reduce(x)

out = self.conv3x3(out)

out = self.increase(out)

out += self.shortcut(x)

return F.relu(out)

class ResLayer(nn.Sequential):

def __init__(self, n_layers, in_channels, out_channels, stride, dilation, multi_grads=None):

super(ResLayer, self).__init__()

if multi_grads is None:

multi_grads = [1 for _ in range(n_layers)]

else:

assert n_layers == len(multi_grads)

for i in range(n_layers):

self.add_module('block{}'.format(i + 1),

Bottlenenck(in_channels=in_channels if i == 0 else out_channels,

out_channels=out_channels,

stride=stride if i == 0 else 1, # 仅在第一个block进行对tensor下采样

dilation=dilation * multi_grads[i],

downsample=True if i == 0 else False)) # 仅在第一个block进行下采样 x

class Stem(nn.Sequential):

def __init__(self, out_channels):

super(Stem, self).__init__()

self.add_module('stem', ConvBnReLUBlock(1, 3, out_channels, 7, 2, 3, 1))

self.add_module('maxpool', nn.MaxPool2d(3, 2, 1, ceil_mode=False))

class ResNetBase(nn.Sequential):

def __init__(self, n_blocks=None, n_rates=None):

super(ResNetBase, self).__init__()

if n_blocks is None:

n_blocks = [3, 4, 23, 3]

else:

assert len(n_blocks) == 4

if n_rates is None:

n_rates = [1, 1, 2, 4]

else:

assert len(n_rates) == 4

chs = [64 * 2 ** k for k in range(6)]

self.add_module('layer0', Stem(chs[0])) # 下采样x4

self.add_module('layer1', ResLayer(n_blocks[0], chs[0], chs[2], 1, n_rates[0]))

self.add_module('layer2', ResLayer(n_blocks[1], chs[2], chs[3], 2, n_rates[1])) # 下采样x2

self.add_module('layer3', ResLayer(n_blocks[2], chs[3], chs[4], 1, n_rates[2])) # resnet-101设置stride=2，下采样x2

self.add_module('layer4', ResLayer(n_blocks[3], chs[4], chs[5], 1, n_rates[3])) # resnet-101设置stride=2，下采样x2

try:

self.load_pretrained_layers()

except:

print("Can not load parameters from pretrained ResNet-101.")

self.init_param()

for c in self.modules():

if isinstance(c, nn.BatchNorm2d):

for p in c.parameters():

p.requires_grad = False

def init_param(self):

for n, c in self.named_parameters():

if 'weight' in n and 'bn' not in n:

nn.init.xavier_normal_(c)

if 'bias' in n:

nn.init.constant_(c, 0)

if 'bn' in n:

nn.init.constant_(c, 1)

def load_pretrained_layers(self):

pretrained_net = torchvision.models.resnet101(pretrained=True)

pretrained_state_dict = pretrained_net.state_dict()

state_dict = self.state_dict()

for layer, weight in zip(state_dict.keys(), pretrained_state_dict.values()):

assert state_dict[layer].shape == weight.shape

state_dict[layer] = weight

self.load_state_dict(state_dict)

3.2 ASPP

class ASPP(nn.Module):

def __init__(self, in_channels, out_channels, rates):

super(ASPP, self).__init__()

for i, rate in enumerate(rates):

self.add_module('aspp{}'.format(i),

nn.Conv2d(in_channels, out_channels, 3, 1, rate, rate, bias=True))

self.init_param()

def init_param(self):

for c in self.children():

nn.init.xavier_normal_(c.weight)

nn.init.constant_(c.bias, 0)

def forward(self, x):

return sum([c(x) for c in self.children()])

3.3 Poly Learning Rate Policy

def poly_lr_scheduler(optimizer, n_iter, lr_decay_iter=1, max_iter=100, power=0.9):

"""

与作者给出的按iter衰减不同，我对epoch进行衰减，最大100个epoch，直接传入optimizer进行lr衰减

"""

if n_iter % lr_decay_iter == 0 and n_iter <= max_iter:

for param_gourp in optimizer.param_groups:

param_gourp['lr'] *= (1 - n_iter / max_iter) ** power

print("DECAYING learning rate... The new LR is %f\n" % (optimizer.param_groups[0]['lr']))

3.4Multiscale Processing

class MultiscaleProcess(nn.Module):

def __init__(self, net, scales=None):

super(MultiscaleProcess, self).__init__()

self.net = net

if scales is None:

scales = [0.5, 0.75]

self.scales = scales

def forward(self, x):

out = self.net(x)

_, _, h, w = out.shape

out_mutiscales = []

for s in self.scales:

img = F.interpolate(x, scale_factor=s, mode='bilinear', align_corners=False)

out_mutiscales.append(self.net(img))

out_all = [out] + [F.interpolate(s, size=(h, w), mode='bilinear', align_corners=False) for s in out_mutiscales]

out_max = torch.stack(out_all, dim=0).max(dim=0)[0]

if self.net.training:

return [out] + out_mutiscales + [out_max]

else:

return out_max

3.5 坑点

我主要试验了如下的几种组合：

ResNet-101 + ASPP ：仅使用VOC训练，结果一直都是黑的，因此将进行了特征提取层和打分层分开进行如下的试验测试，看看ResNet和ASPP模型是不是哪里出问题了

VGG-16 + ASPP ：很容易仅使用VOC训练出了50+的IoU，大概1个epoch就行

ResNet-101 + LargeFOV ：仅使用VOC较难训练，20个epoch才能达到接近50的mIoU，但mIoU基本收敛在40+实际分割效果并不是很好，且对于部分类别的效果不佳

试验后发现两部分模型都没有问题，且较之VGG-16和LargeFOV更难训练，改用VOC AUG进行训练

ResNet-101 + ASPP + Poly Learning Rate Policy ：使用VOC AUG进行训练，十几个epoch后就能训练出60+的IoU，各个类别的分割效果都很好，Poly Learning Rate Policy的确能够提升mIoU

ResNet-101 + ASPP + Poly Learning Rate Policy + Multiscale Processing：使用VOC AUG进行训练，能够更快速的收敛

如上两个结构的模型结构的模型我都没有训练太多epoch(也就20刚出头，算力有限😓)，但也确实证明了模型的优秀能力。

本次跑模型最大的坑点在于预训练的BN层：使用预训练模型的参数时需要设置固定BN层的参数不被优化，不然会带来问题。对于BN层，在训练时，是对mini-batch的训练数据进行归一化，也即用每一批数据的均值和方差，对于我设置的512 x 512的训练图像，只能设置batch=4，给模型带来了严重问题。

此外，尽管模型是基于预训练的ResNet-101网络进行进行微调，但其参数远比VGG-16多太多了，仅使用VOC数据集进行训练很难达到预期效果，所以必须使用VOC AUG数据集进行训练才能避免“全黑”分割结果的出现。

在进行Multiscale Processing时，需要对ground truth label同样进行下采样，于是我使用了如下操做：

# 导致训练失败的问题代码

import torchvision.transforms.functional as F

for pred in preds:

_, _, h, w = pred.shape

new_label = []

for l in label_img:

new_l=F.to_tensor(F.resize(F.to_pil_image(l.float().cpu()), (h, w), Image.NEAREST)).long().to(device)

new_label.append(label_img)

label_image = torch.LongTensor(VOC_COLORMAP)[new_label[-1].flatten()].reshape(h, w, 3)

由于使用Nearest插值方法的限制，torch.nn.functional.interploate并不支持Nearest，所以只能ground truth转换为按图片并进行插值。

F.to_pil_image()需要传入进行归一化处理后的FloatTensor，但我传入的数据是有问题的，因此导致了训练中的说莫名的错误；类似的F.to_tensor()也会将输入进行归一化，所以以上的操作达不到我的预期目标。因此我新写了一个函数来将ground truth label进行rescale。

# 可以完成训练

def ground_truth_to_resized_label(label_img, size):

l = label_img.cpu().numpy().astype(np.uint8)

l = Image.fromarray(l).resize(size, resample=Image.NEAREST)

l = torch.LongTensor(np.array(l))

return l

基本上DenseCRF后处理对于DeepLabv2模型的提升已经不是十分明显了，甚至可能其到反效果，现在基本上只有比赛刷分还能用到，如果对于DenseCRF后处理感兴趣可参考DeepLabv1中相关的介绍和后处理代码，这里由于算力的限制，就不后处理啦😂

结果

Reference

[1] Chen, L. C. , Papandreou, G. , Kokkinos, I. , Murphy, K. , & Yuille, A. L. . (2017). Deeplab: semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. IEEE Transactions on Pattern Analysis & Machine Intelligence, 1-1.

转载请说明出处。