mobilenetv2语义分割代码

前言

FCN模型结构往简答说，就是先用VGG类似的卷积网络进行特征提取，然后再对特征图进行反卷积(deconvolution)来将其投影到像素空间从而实现逐个逐个像素的分类。 

结合代码分析模型结构

说起来简单，不过还是需要结合代码把其细节部分讲清楚。

在train.py中，只用了两行代码就创建了FCN模型。

vgg_model = models.VGGNet(requires_grad=True)
fcn_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)

这两行代码的实现部分在models.py中，它是先创建vgg模型，然后基于它再创建FCN模型。

我们进去models.py先看看class VGGNet的构造函数。

def __init__(self, pretrained=True, model='vgg16', requires_grad=True, remove_fc=True, show_params=False):
        super().__init__(make_layers(cfg[model]))
        self.ranges = ranges[model]

        if pretrained:
            vgg16 = models.vgg16(pretrained=False)
            vgg16.load_state_dict(torch.load('/home/xxx/models/vgg16-397923af.pth'))
            # exec("self.load_state_dict(models.%s(pretrained=True).state_dict())" % model)

        if not requires_grad:
            for param in super().parameters():
                param.requires_grad = False

        if remove_fc: # delete redundant fully-connected layer params, can save memory
            del self.classifier

        if show_params:
            for name, param in self.named_parameters():
                print(name, param.size())

这段代码有几个注意点：

1） 参数model='vgg16'  因为 VGG模型有VGG11, VGG13, VGG16以及VGG19等类型，这里使用model缺省参数VGG16。

2）通过cfg[model]就可以获得对应VGG模型配置文件，然后通过make_layers()就把VGG模型建立起来。

3）因为pretrained参数为True，所以需要加载一个预训练好的模型放到本地目录。因为该模型比较大，所以最好先下载好。

4）remove_fc为True，因为FCN是全卷积网络模型，所以需要把VGG最后的全连接层去掉。

VGG提取特征后，FCN模型就开始对特征图使用deconvolution进行upsample，按不同的upsample方式，FCN可以分为FCN8s，FCN16s以及FCN32s。 如下图所示。

首先要搞清楚的一点，就是VGG16一共有5个block，每个block的最后一层就是MaxPooling，对应到上图就是pool1,pool2,pool3,pool4以及pool5。

FCN32s就是直接对pool5进行32倍的放大（相当于stride为32的反卷积）变回输入图像分辨率大小的像素预测空间；FCN16s则是pool5的2倍放大+pool4，然后再16倍放大到输入图像大小的空间；FCN8s则是pool3加上(pool4+pool5的2倍upsample)的2倍upsample，再8倍upsample到输入图像大小。 从这里可以看出，FCN32s最为简单粗暴，所以准确度相对最低，而FCN8s考虑了3个pool层并融合，所以相对较为精细。

下面是FCN8s的forward函数，它在train和predict时被调用：

def forward(self, x):
        output = self.pretrained_net.forward(x)
        x5 = output['x5']  # size=[n, 512, x.h/32, x.w/32]
        x4 = output['x4']  # size=[n, 512, x.h/16, x.w/16]
        x3 = output['x3']  # size=[n, 512, x.h/8, x.w/8]

        score = self.relu(self.deconv1(x5))                  # size=[n, 512, x.h/16, x.w/16]
        score = self.bn1(score + x4)                         # element-wise add, size=[n, 512, x.h/16, x.w/16]
        score = self.relu(self.deconv2(score))               # size=[n, 256, x.h/8, x.w/8]
        score = self.bn2(score+x3)
        score = self.bn3(self.relu(self.deconv3(score)))     # size=[n, 128, x.h/4, x.w/4]
        score = self.bn4(self.relu(self.deconv4(score)))     # size=[n, 64, x.h/2, x.w/2]
        score = self.bn5(self.relu(self.deconv5(score)))     # size=[n, 32, x.h, x.w]
        score = self.classifier(score)                       # size=[n, n_class, x.h, x.w]

        return score

 上面代码中，x3，x4和x5分别是pool3，pool4以及pool5的动态输出。然后对x5进行2倍upsample，并和x4进行逐像素相加，其和再做一个2倍upsample，最后和pool3逐像素相加，其和连续做3次2倍upsample就得到和输入图像分辨率相等像素预测值矩阵。