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卷积设计的核心理念是临近的像素携带的信息相似。

该图片来自：vdumoulin/conv_arithmetic

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（lr\ *=\ (1-\frac{iter}{max\_iter})^{power},power=0.9）能够得到更好的表现。

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.

[2] kazuto1011/ deeplab-pytorch

[3] doiken23/DeepLab_pytorch

转载请说明出处。

最后编辑于 ：
©著作权归作者所有,转载或内容合作请联系作者
  • 人面猴
    序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
    沈念sama阅读 215,723评论 6赞 498
  • 死咒
    序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
    沈念sama阅读 92,003评论 3赞 391
  • 救了他两次的神仙让他今天三更去死
    文/潘晓璐 我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
    开封第一讲书人阅读 161,512评论 0赞 351
  • 道士缉凶录：失踪的卖姜人
    文/不坏的土叔 我叫张陵，是天一观的道长。 经常有香客问我，道长，这世上最难降的妖魔是什么？ 我笑而不...
    开封第一讲书人阅读 57,825评论 1赞 290
  • 港岛之恋（遗憾婚礼）
    正文 为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
    茶点故事阅读 66,874评论 6赞 388
  • 恶毒庶女顶嫁案：这布局不是一般人想出来的
    文/花漫 我一把揭开白布。 她就那样静静地躺着，像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上，一...
    开封第一讲书人阅读 50,841评论 1赞 295
  • 城市分裂传说
    那天，我揣着相机与录音，去河边找鬼。 笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。 我是一名探鬼主播，决...
    沈念sama阅读 39,812评论 3赞 416
  • 双鸳鸯连环套：你想象不到人心有多黑
    文/苍兰香墨 我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
    开封第一讲书人阅读 38,582评论 0赞 271
  • 万荣杀人案实录
    序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
    沈念sama阅读 45,033评论 1赞 308
  • 护林员之死
    正文 独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,309评论 2赞 331
  • 白月光启示录
    正文 我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,450评论 1赞 345
  • 活死人
    序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
    沈念sama阅读 35,158评论 5赞 341
  • 日本核电站爆炸内幕
    正文 年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
    茶点故事阅读 40,789评论 3赞 325
  • 男人毒药：我在死后第九天来索命
    文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,409评论 0赞 21
  • 一桩弑父案，背后竟有这般阴谋
    文/苍兰香墨 我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,609评论 1赞 268
  • 情欲美人皮
    我被黑心中介骗来泰国打工， 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。 一个月前我还...
    沈念sama阅读 47,440评论 2赞 368
  • 代替公主和亲
    正文 我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
    茶点故事阅读 44,357评论 2赞 352