Machine World | Math World

记录实验过程

【经典复现】语义分割经典-FCN(Backbone:VGG)

【背景】

        无聊看了几篇论/博文,有必要记录一下笔记吧。

【环境】

  • CUDA: 9.2.148

  • Torch:1.2.0

  • OS: Ubuntu 16.04

  • HW: Nvidia Tesla P100 / Nvidia GTX 1080Ti / Nvidia RTX 2080Ti

【一句话概括FCN】

    FCN将传统卷积网络后面的全连接层换成了卷积层,这样网络输出不再是类别而是 heatmap;同时为了解决因为卷积和池化对图像尺寸的影响,提出使用上采样的方式恢复。其结构图如下:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

    本文重点在于复现,复现FCN分为四步详解,数据集处理,选取Backbone,复现FCN框架,训练方法。接下一一阐述。

【数据集篇】

数据准备:

    手提包数据集:点击下载

使用库:

  • cv2: 用以对数据集图片的resize操作,注:cv的颜色通道与PIL有区别。

  • torch.utils.data.Dataloader, Dataset, random_split :构造数据集用

  • torchvision.transforms: 数据集预处理

  • numpy: onehot操作

数据Demo:

Training Data:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

Training Label:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

    从这个手提包数据集可以看出,这是个二分类的,就是只分割出手提包 和 背景 两个类别。所以label处黑色的表示手提包,白色的就是无关的背景。

onehot化:将Label化为一个一维向量(类别映射)

def onehot(data, n):
    buf = np.zeros(data.shape + (n, )) #设data的shape为(a,b),构造一个zeros矩阵,shape为(a, b, n)
    nmsk = np.arange(data.size) * n + data.ravel() 
    buf.ravel()[nmsk-1] = 1
    return buf

但其实,pytorch还给出了更便捷的方法:

class_num = 10
batch_size = 4
label = torch.LongTensor(batch_size, 1).random_() % class_num
print(label.size()) 
one_hot = torch.zeros(batch_size, class_num).scatter_(1, label, 1)
print(one_hot)

label的形状必须是[n,1]的,也就是必须是二维的,且第二个维度长度为1,如果是一维度的,则需要升维度,代码如下:

class_num = 10
batch_size = 4
label = torch.LongTensor(batch_size).random_() % class_num
print(label.size())
label = torch.unsqueeze(label,dim=1)
print(label.size())

假设已知数据集中类别(以VOC 2012为例),可直接做映射:

classes = ['background','aeroplane','bicycle','bird','boat',
           'bottle','bus','car','cat','chair','cow','diningtable',
           'dog','horse','motorbike','person','potted plant',
           'sheep','sofa','train','tv/monitor']
# RGB color for each 
classcolormap = [[0,0,0],[128,0,0],[0,128,0], [128,128,0], [0,0,128],
            [128,0,128],[0,128,128],[128,128,128],[64,0,0],[192,0,0],
            [64,128,0],[192,128,0],[64,0,128],[192,0,128],
            [64,128,128],[192,128,128],[0,64,0],[128,64,0],
            [0,192,0],[128,192,0],[0,64,128]]

映射方法:

cm2lbl = np.zeros(256**3) # 每个像素点有 0 ~ 255 的选择,RGB 三个通道
for i,cm in enumerate(colormap):
    cm2lbl[(cm[0]*256+cm[1])*256+cm[2]] = i 
# 建立索引
def image2label(im):
    data = np.array(im, dtype='int32')
    idx = (data[:, :, 0] * 256 + data[:, :, 1]) * 256 + data[:, :, 2]
    return np.array(cm2lbl[idx], dtype='int64') # 根据索引得到 label 矩阵

理解完label映射后,开始构造数据集:(重写Dataloader类)

transform = transforms.Compose([
    transforms.ToTensor(), 
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

class BagDataset(Dataset):

    def __init__(self, transform=None):
        self.transform = transform
        
    def __len__(self):
        return len(os.listdir('bag_data'))

    def __getitem__(self, idx):
        img_name = os.listdir('bag_data')[idx]
        imgA = cv2.imread('bag_data/'+img_name)
        imgA = cv2.resize(imgA, (160, 160))
        imgB = cv2.imread('bag_data_msk/'+img_name, 0)
        imgB = cv2.resize(imgB, (160, 160))
        imgB = imgB/255
        imgB = imgB.astype('uint8')
        imgB = onehot(imgB, 2)
        imgB = imgB.transpose(2,0,1)
        imgB = torch.FloatTensor(imgB)
        #print(imgB.shape)
        if self.transform:
            imgA = self.transform(imgA)    

        return imgA, imgB

【网络篇-backbone】

VGG:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

结构:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

Network定义如下:

class VGGNet(VGG):
    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:
            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

        # delete redundant fully-connected layer params, can save memory
        # 去掉vgg最后的全连接层(classifier)
        if remove_fc:  
            del self.classifier

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

    def forward(self, x):
        output = {}
        # get the output of each maxpooling layer (5 maxpool in VGG net)
        for idx, (begin, end) in enumerate(self.ranges):
        #self.ranges = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)) (vgg16 examples)
            for layer in range(begin, end):
                x = self.features[layer](x)
            output["x%d"%(idx+1)] = x

        return output


ranges = {
    'vgg11': ((0, 3), (3, 6),  (6, 11),  (11, 16), (16, 21)),
    'vgg13': ((0, 5), (5, 10), (10, 15), (15, 20), (20, 25)),
    'vgg16': ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)),
    'vgg19': ((0, 5), (5, 10), (10, 19), (19, 28), (28, 37))
}

# Vgg-Net config 
# Vgg网络结构配置
cfg = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

# make layers using Vgg-Net config(cfg)
# 由cfg构建vgg-Net
def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)
 '''
VGG-16网络参数

Sequential(
  (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU(inplace)
  (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (3): ReLU(inplace)
  (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (6): ReLU(inplace)
  (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (8): ReLU(inplace)
  (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (11): ReLU(inplace)
  (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (13): ReLU(inplace)
  (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (15): ReLU(inplace)
  (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (18): ReLU(inplace)
  (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (20): ReLU(inplace)
  (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (22): ReLU(inplace)
  (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (25): ReLU(inplace)
  (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (27): ReLU(inplace)
  (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (29): ReLU(inplace)
  (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

'''

【网络篇-FCN】

在FCN中核心就是将传统的卷积神经网络backbone最后一层替换成全卷积层。如下图:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

图中,image是原图像,conv1,conv2..,conv5为卷积操作,pool1,pool2,..pool5为pool操作(pool就是使得图片变为原图的1/2),注意con6-7是最后的卷积层,最右边一列是upsample后的end to end结果。必须说明的是图中nx是指对应的特征图上采样n倍(即变大n倍),并不是指有n个特征图,如32x upsampled 中的32x是图像只变大32倍,不是有32个上采样图像,又如2x conv7是指conv7的特征图变大2倍。

(1)FCN-32s过程

只需要留意第一行,网络里面有5个pool,所以conv7的特征图是原始图像1/32,可以发现最左边image的是32×32(假设以倍数计),同时我们知道在FCN中的卷积是不会改变图像大小(或者只有少量像素的减少,特征图大小基本不会小很多)。看到pool1是16×16,pool2是8×8,pool3是4×4,pool4是2×2,pool5是1×1,所以conv7对应特征图大小为1×1,然后再经过32x upsampled prediction 图片变回32×32。FCN作者在这里增加一个卷积层,卷积后的大小为输入图像的32(2^5)倍,我们简单假设这个卷积核大小也为32,这样就是需要通过反馈训练32×32个权重变量即可让图像实现end to end,完成了一个32s的upsample。FCN作者称做后卷积,他也提及可以称为反卷积。事实上在源码中卷积核的大小为64,同时没有偏置bias。还有一点就是FCN论文中最后结果都是21×*,这里的21是指FCN使用的数据集分类,总共有21类。 

(2)FCN-16s过程

现在我们把1,2两行一起看,忽略32x upsampled prediction,说明FCN-16s的upsample过程。FCN作者在conv7先进行一个2x conv7操作,其实这里也只是增加1个卷积层,这次卷积后特征图的大小为conv7的2倍,可以从pool5与2x conv7中看出来。此时2x conv7与pool4的大小是一样的,FCN作者提出对pool4与2x conv7进行一个fuse操作(事实上就是将pool4与2x conv7相加,另一篇博客说是拼接,个人认为是拼接)。fuse结果进行16x upsampled prediction,与FCN-32s一样,也是增加一个卷积层,卷积后的大小为输入图像的16(2^4)倍。我们知道pool4的大小是2×2,放大16倍,就是32×32,这样最后图像大小也变为原来的大小,至此完成了一个16s的upsample。现在我们可以知道,FCN中的upsample实际是通过增加卷积层,通过bp反馈的训练方法训练卷积层达到end to end,这时卷积层的作用可以看作是pool的逆过程。 

(3)FCN-8s过程

这是我们看第1行与第3行,忽略32x upsampled prediction。conv7经过一次4x upsample,即使用一个卷积层,特征图输出大小为conv7的4倍,所得4x conv7的大小为4×4。然后pool4需要一次2x upsample,变成2x pool4,大小也为4×4。再把4x conv7,2x pool4与pool3进行fuse,得到求和后的特征图。最后增加一个卷积层,使得输出图片大小为pool3的8倍,也就是8x upsampled prediction的过程,得到一个end to end的图像。实验表明FCN-8s优于FCN-16s,FCN-32s。 

我们可以发现,如果继续仿照FCN作者的步骤,我们可以对pool2,pool1实现同样的方法,可以有FCN-4s,FCN-2s,最后得到end to end的输出。这里作者给出了明确的结论,超过FCN-8s之后,结果并不能继续优化。

结合上述的FCN的全卷积与upsample,在upsample最后加上softmax,就可以对不同类别的大小概率进行估计,实现end to end。最后输出的图是一个概率估计,对应像素点的值越大,其像素为该类的结果也越大。FCN的核心贡献在于提出使用卷积层通过学习让图片实现end to end分类。

事实上,FCN有一些短处,例如使用了较浅层的特征,因为fuse操作会加上较上层的pool特征值,导致高维特征不能很好得以使用,同时也因为使用较上层的pool特征值,导致FCN对图像大小变化有所要求,如果测试集的图像远大于或小于训练集的图像,FCN的效果就会变差。

代码如下:

class FCNs(nn.Module):

    def __init__(self, pretrained_net, n_class):
        super().__init__()
        self.n_class = n_class
        self.pretrained_net = pretrained_net
        self.relu    = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn1     = nn.BatchNorm2d(512)
        self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn2     = nn.BatchNorm2d(256)
        self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn3     = nn.BatchNorm2d(128)
        self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn4     = nn.BatchNorm2d(64)
        self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn5     = nn.BatchNorm2d(32)
        self.classifier = nn.Conv2d(32, n_class, kernel_size=1)
        # classifier is 1x1 conv, to reduce channels from 32 to n_class

    def forward(self, x):
        output = self.pretrained_net(x)
        x5 = output['x5']
        x4 = output['x4']
        x3 = output['x3']
        x2 = output['x2']
        x1 = output['x1']
        score = self.bn1(self.relu(self.deconv1(x5)))
        score = score + x4
        score = self.bn2(self.relu(self.deconv2(score)))
        score = score + x3
        score = self.bn3(self.relu(self.deconv3(score)))  
        score = score + x2
        score = self.bn4(self.relu(self.deconv4(score)))
        score = score + x1
        score = self.bn5(self.relu(self.deconv5(score)))
        score = self.classifier(score)
        return score

【训练篇】

训练分四阶段:

第一阶段:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

    以经典的分类网络为初始化。最后两级是全连接(红色),参数弃去不用。

第二阶段:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

    从特征小图(16*16*4096)预测分割小图(16*16*21),之后直接升采样为大图。 

    反卷积(橙色)的步长为32,这个网络称为FCN-32s。 

第三阶段:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

    上采样采样分为两次完成(橙色×2)。 

    在第二次上采样前,把第4个pooling层(绿色)的预测结果(蓝色)融合进来。使用跳级结构提升精确性。 

    第二次反卷积步长为16,这个网络称为FCN-16s。 

第四阶段:

《【经典复现】语义分割经典-FCN(Backbone:VGG)》

    上采样分为三次完成(橙色×3)。 

    进一步融合了第3个pooling层的预测结果。 

    第三次反卷积步长为8,记为FCN-8s。 

但本文融合了五层polling的预测结果。

训练Code:

from datetime import datetime

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import visdom

from FCN import FCN8s, FCN16s, FCN32s, FCNs, VGGNet


def train(epo_num=50, show_vgg_params=False):
    #初始化数据集
    train_size = int(0.9 * len(bag))
    test_size = len(bag) - train_size
    train_dataset, test_dataset = random_split(bag, [train_size, test_size])

    train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
    test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=4)
    #开启visdom
    vis = visdom.Visdom()
    #GPU计算
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    vgg_model = VGGNet(requires_grad=True, show_params=show_vgg_params)
    fcn_model = FCNs(pretrained_net=vgg_model, n_class=2)
    fcn_model = fcn_model.to(device)
    criterion = nn.BCELoss().to(device) #损失函数
    optimizer = optim.SGD(fcn_model.parameters(), lr=1e-2, momentum=0.7) #优化器(SGD)

    all_train_iter_loss = []
    all_test_iter_loss = []

    # start timing
    prev_time = datetime.now()
    for epo in range(epo_num):
        
        train_loss = 0
        fcn_model.train()
        for index, (bag, bag_msk) in enumerate(train_dataloader):
            # bag.shape is torch.Size([4, 3, 160, 160])
            # bag_msk.shape is torch.Size([4, 2, 160, 160])

            bag = bag.to(device)
            bag_msk = bag_msk.to(device)

            optimizer.zero_grad()
            output = fcn_model(bag)
            output = torch.sigmoid(output) # output.shape is torch.Size([4, 2, 160, 160])
            loss = criterion(output, bag_msk)
            loss.backward()
            iter_loss = loss.item()
            all_train_iter_loss.append(iter_loss)
            train_loss += iter_loss
            optimizer.step()

            output_np = output.cpu().detach().numpy().copy() # output_np.shape = (4, 2, 160, 160)  
            output_np = np.argmin(output_np, axis=1)
            bag_msk_np = bag_msk.cpu().detach().numpy().copy() # bag_msk_np.shape = (4, 2, 160, 160) 
            bag_msk_np = np.argmin(bag_msk_np, axis=1)

            if np.mod(index, 15) == 0:
                print('epoch {}, {}/{},train loss is {}'.format(epo, index, len(train_dataloader), iter_loss))
                # vis.close()
                vis.images(output_np[:, None, :, :], win='train_pred', opts=dict(title='train prediction')) 
                vis.images(bag_msk_np[:, None, :, :], win='train_label', opts=dict(title='label'))
                vis.line(all_train_iter_loss, win='train_iter_loss',opts=dict(title='train iter loss'))

            # plt.subplot(1, 2, 1) 
            # plt.imshow(np.squeeze(bag_msk_np[0, ...]), 'gray')
            # plt.subplot(1, 2, 2) 
            # plt.imshow(np.squeeze(output_np[0, ...]), 'gray')
            # plt.pause(0.5)

        
        test_loss = 0
        fcn_model.eval()
        with torch.no_grad():
            for index, (bag, bag_msk) in enumerate(test_dataloader):

                bag = bag.to(device)
                bag_msk = bag_msk.to(device)

                optimizer.zero_grad()
                output = fcn_model(bag)
                output = torch.sigmoid(output) # output.shape is torch.Size([4, 2, 160, 160])
                loss = criterion(output, bag_msk)
                iter_loss = loss.item()
                all_test_iter_loss.append(iter_loss)
                test_loss += iter_loss

                output_np = output.cpu().detach().numpy().copy() # output_np.shape = (4, 2, 160, 160)  
                output_np = np.argmin(output_np, axis=1)
                bag_msk_np = bag_msk.cpu().detach().numpy().copy() # bag_msk_np.shape = (4, 2, 160, 160) 
                bag_msk_np = np.argmin(bag_msk_np, axis=1)
        
                if np.mod(index, 15) == 0:
                    print(r'Testing... Open http://localhost:8097/ to see test result.')
                    # vis.close()
                    vis.images(output_np[:, None, :, :], win='test_pred', opts=dict(title='test prediction')) 
                    vis.images(bag_msk_np[:, None, :, :], win='test_label', opts=dict(title='label'))
                    vis.line(all_test_iter_loss, win='test_iter_loss', opts=dict(title='test iter loss'))

                # plt.subplot(1, 2, 1) 
                # plt.imshow(np.squeeze(bag_msk_np[0, ...]), 'gray')
                # plt.subplot(1, 2, 2) 
                # plt.imshow(np.squeeze(output_np[0, ...]), 'gray')
                # plt.pause(0.5)


        cur_time = datetime.now()
        h, remainder = divmod((cur_time - prev_time).seconds, 3600)
        m, s = divmod(remainder, 60)
        time_str = "Time %02d:%02d:%02d" % (h, m, s)
        prev_time = cur_time

        print('epoch train loss = %f, epoch test loss = %f, %s'
                %(train_loss/len(train_dataloader), test_loss/len(test_dataloader), time_str))
        

        if np.mod(epo, 5) == 0:
            torch.save(fcn_model, 'checkpoints/fcn_model_{}.pt'.format(epo))
            print('saveing checkpoints/fcn_model_{}.pt'.format(epo))


if __name__ == "__main__":
  train(epo_num=100, show_vgg_params=True)

【总结】

    FCN 的优势在于: 

        – 可以接受任意大小的输入图像(没有全连接层) 

        – 更加高效,避免了使用邻域带来的重复计算和空间浪费的问题。

    其不足也很突出: 

        – 得到的结果还不够精细 。进行8倍上采样虽然比32倍的效果好了很多,但是上采样的结果还是比较模糊和平滑,对图像中的细节不敏感。

        -是对各个像素进行分类,没有充分考虑像素与像素之间的关系。忽略了在通常的基于像素分类的分割方法中使用的空间规整(spatial regularization)步骤,缺乏空间一致性。

【参考文献】

点赞

发表评论

电子邮件地址不会被公开。 必填项已用*标注