全卷积网络（FCN）实战：使用FCN实现语义分割

摘要：FCN对图像进行像素级的分类，从而解决了语义级别的图像分割问题。

本文分享自华为云社区《全卷积网络（FCN）实战：使用FCN实现语义分割》，作者： AI浩。

FCN对图像进行像素级的分类，从而解决了语义级别的图像分割（semantic segmentation）问题。与经典的CNN在卷积层之后使用全连接层得到固定长度的特征向量进行分类（全联接层＋softmax输出）不同，FCN可以接受任意尺寸的输入图像，采用反卷积层对最后一个卷积层的feature map进行上采样, 使它恢复到输入图像相同的尺寸，从而可以对每个像素都产生了一个预测, 同时保留了原始输入图像中的空间信息, 最后在上采样的特征图上进行逐像素分类。

下图是语义分割所采用的全卷积网络(FCN)的结构示意图：

传统的基于CNN的分割方法缺点？

传统的基于CNN的分割方法：为了对一个像素分类，使用该像素周围的一个图像块作为CNN的输入，用于训练与预测，这种方法主要有几个缺点：

1）存储开销大，例如，对每个像素使用15 * 15的图像块，然后不断滑动窗口，将图像块输入到CNN中进行类别判断，因此，需要的存储空间随滑动窗口的次数和大小急剧上升；

2）效率低下，相邻像素块基本上是重复的，针对每个像素块逐个计算卷积，这种计算有很大程度上的重复；

3）像素块的大小限制了感受区域的大小，通常像素块的大小比整幅图像的大小小很多，只能提取一些局部特征，从而导致分类性能受到限制。
而全卷积网络(FCN)则是从抽象的特征中恢复出每个像素所属的类别。即从图像级别的分类进一步延伸到像素级别的分类。

FCN改变了什么?

对于一般的分类CNN网络，如VGG和Resnet，都会在网络的最后加入一些全连接层，经过softmax后就可以获得类别概率信息。但是这个概率信息是1维的，即只能标识整个图片的类别，不能标识每个像素点的类别，所以这种全连接方法不适用于图像分割。

而FCN提出可以把后面几个全连接都换成卷积，这样就可以获得一张2维的feature map，后接softmax层获得每个像素点的分类信息，从而解决了分割问题，如图。

FCN缺点

（1）得到的结果还是不够精细。进行8倍上采样虽然比32倍的效果好了很多，但是上采样的结果还是比较模糊和平滑，对图像中的细节不敏感。
（2）对各个像素进行分类，没有充分考虑像素与像素之间的关系。忽略了在通常的基于像素分类的分割方法中使用的空间规整（spatial regularization）步骤，缺乏空间一致性。

数据集

本例的数据集采用PASCAL VOC 2012 数据集，它有二十个类别：

Person:person

Animal: bird, cat, cow, dog, horse, sheep

Vehicle:aeroplane, bicycle, boat, bus, car, motorbike, train

Indoor: bottle, chair, dining table, potted plant, sofa, tv/monitor

下载地址：deep-learning-for-image-processing/pytorch_segmentation/fcn at master · WZMIAOMIAO/deep-learning-for-image-processing (github.com)

其他的代码也有很多，这篇比较好理解！

其实还有个比较好的图像分割库：https://github.com/qubvel/segmentation_models.pytorch

这个图像分割集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造。在后面的文章，我也会使用这个库演示。

项目结构

├── src: 模型的backbone以及FCN的搭建
├── train_utils: 训练、验证以及多GPU训练相关模块
├── my_dataset.py: 自定义dataset用于读取VOC数据集
├── train.py: 以fcn_resnet50(这里使用了Dilated/Atrous Convolution)进行训练
├── predict.py: 简易的预测脚本，使用训练好的权重进行预测测试
├── validation.py: 利用训练好的权重验证/测试数据的mIoU等指标，并生成record_mAP.txt文件
└── pascal_voc_classes.json: pascal_voc标签文件

由于代码很多不能一一讲解，所以，接下来对重要的代码做剖析。

自定义数据集读取

my_dataset.py自定义数据读取的方法，代码如下：

import os
import torch.utils.data as data
from PIL import Image

class VOCSegmentation(data.Dataset):
    def __init__(self, voc_root, year="2012", transforms=None, txt_name: str = "train.txt"):
        super(VOCSegmentation, self).__init__()
        assert year in ["2007", "2012"], "year must be in ['2007', '2012']"
        root = os.path.join(voc_root, "VOCdevkit", f"VOC{year}")
        root=root.replace('\\','/')
        assert os.path.exists(root), "path '{}' does not exist.".format(root)
        image_dir = os.path.join(root, 'JPEGImages')
        mask_dir = os.path.join(root, 'SegmentationClass')

        txt_path = os.path.join(root, "ImageSets", "Segmentation", txt_name)
        txt_path=txt_path.replace('\\','/')
        assert os.path.exists(txt_path), "file '{}' does not exist.".format(txt_path)
        with open(os.path.join(txt_path), "r") as f:
            file_names = [x.strip() for x in f.readlines() if len(x.strip()) > 0]

        self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names]
        self.masks = [os.path.join(mask_dir, x + ".png") for x in file_names]
        assert (len(self.images) == len(self.masks))
        self.transforms = transforms

导入需要的包。

定义VOC数据集读取类VOCSegmentation。在init方法中，核心是读取image列表和mask列表。

  def __getitem__(self, index):
        img = Image.open(self.images[index]).convert('RGB')
        target = Image.open(self.masks[index])

        if self.transforms is not None:
            img, target = self.transforms(img, target)
        return img, target

__getitem__方法是获取单张图片和图片对应的mask，然后对其做数据增强。

 def collate_fn(batch):
        images, targets = list(zip(*batch))
        batched_imgs = cat_list(images, fill_value=0)
        batched_targets = cat_list(targets, fill_value=255)
        return batched_imgs, batched_targets

collate_fn方法是对一个batch中数据调用cat_list做数据对齐。

在train.py中torch.utils.data.DataLoader调用

 train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size,
                                               num_workers=num_workers,
                                               shuffle=True,
                                               pin_memory=True,
                                               collate_fn=train_dataset.collate_fn)
  val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=1,
                                             num_workers=num_workers,
                                             pin_memory=True,
                                             collate_fn=val_dataset.collate_fn)

训练

重要参数

打开train.py,我们先认识一下重要的参数：

def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description="pytorch fcn training")
    # 数据集的根目录(VOCdevkit)所在的文件夹
    parser.add_argument("--data-path", default="data/", help="VOCdevkit root")
    parser.add_argument("--num-classes", default=20, type=int)
    parser.add_argument("--aux", default=True, type=bool, help="auxilier loss")
    parser.add_argument("--device", default="cuda", help="training device")
    parser.add_argument("-b", "--batch-size", default=32, type=int)
    parser.add_argument("--epochs", default=30, type=int, metavar="N",
                        help="number of total epochs to train")

    parser.add_argument('--lr', default=0.0001, type=float, help='initial learning rate')
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                        help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)',
                        dest='weight_decay')
    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    # 是否使用混合精度训练
    parser.add_argument("--amp", default=False, type=bool,
                        help="Use torch.cuda.amp for mixed precision training")

    args = parser.parse_args()

    return args

data-path：定义数据集的根目录(VOCdevkit)所在的文件夹

num-classes：检测目标类别数(不包含背景)。

aux：是否使用aux_classifier。

device：使用cpu还是gpu训练，默认是cuda。

batch-size：BatchSize设置。

epochs：epoch的个数。

lr：学习率。

resume：继续训练时候，选择用的模型。

start-epoch：起始的epoch，针对再次训练时，可以不需要从0开始。

amp：是否使用torch的自动混合精度训练。

数据增强

增强调用transforms.py中的方法。

训练集的增强如下：

class SegmentationPresetTrain:
    def __init__(self, base_size, crop_size, hflip_prob=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        # 随机Resize的最小尺寸
        min_size = int(0.5 * base_size)
        # 随机Resize的最大尺寸
        max_size = int(2.0 * base_size)
        # 随机Resize增强。
        trans = [T.RandomResize(min_size, max_size)]
        if hflip_prob > 0:
            #随机水平翻转
            trans.append(T.RandomHorizontalFlip(hflip_prob))
        trans.extend([
            #随机裁剪
            T.RandomCrop(crop_size),
            T.ToTensor(),
            T.Normalize(mean=mean, std=std),
        ])
        self.transforms = T.Compose(trans)

    def __call__(self, img, target):
        return self.transforms(img, target)

训练集增强，包括随机Resize、随机水平翻转、随即裁剪。

验证集增强：

class SegmentationPresetEval:
    def __init__(self, base_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        self.transforms = T.Compose([
            T.RandomResize(base_size, base_size),
            T.ToTensor(),
            T.Normalize(mean=mean, std=std),
        ])

    def __call__(self, img, target):
        return self.transforms(img, target)

验证集的增强比较简单，只有随机Resize。

Main方法

对Main方法，我做了一些修改，修改的代码如下：

 #定义模型，并加载预训练
    model = fcn_resnet50(pretrained=True)
    # 默认classes是21,如果不是21，则要修改类别。
    if num_classes != 21:
        model.classifier[4] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
        model.aux_classifier[4] = torch.nn.Conv2d(256, num_classes, kernel_size=(1, 1), stride=(1, 1))
    print(model)
    model.to(device)
    # 如果有多张显卡，则使用多张显卡
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = torch.nn.DataParallel(model)

模型，我改为pytorch官方的模型了，如果能使用官方的模型尽量使用官方的模型。

默认类别是21，如果不是21，则要修改类别。

检测系统中是否有多张卡，如果有多张卡则使用多张卡不能浪费资源。

如果不想使用所有的卡，而是指定其中的几张卡，可以使用：

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

也可以在DataParallel方法中设定：

model = torch.nn.DataParallel(model,device_ids=[0,1])

如果使用了多显卡，点击关注，第一时间了解华为云新鲜技术~