Pytorch实战语义分割（VOC2012）

杨利霞

Pytorch实战语义分割（VOC2012）
: J4 w: X- e7 o. o, t. q本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。

+ w, C) g3 Q) [2 L: _$ l) ?! t3 z
语义分割中图像有关狗、猫和背景的标签
" R: J8 ^- r1 ]+ R1 ^文章目录
7 G3 @3 |( z5 [2 G6 V
+ {, j1 G) n% \+ ?2 {6 c1 ~8 N7 T& C1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
2.1 导入模块
2.2 下载数据集
2.3 可视化数据
2.4 预处理数据
3 v1 K6 W, d/ O' i! [3 自定义数据集类
3.1 数据集类
0 y$ o; c9 n& x3 }2 d3.2 读取数据集
4 构造模型
4.1 预训练模型
4.2 修改成FCN
4.3 初始化转置卷积层
" g% A9 z  p7 T& Q$ A5 _5 训练模型
7 `; T  a! y/ ~# k6 测试模型
6.1 通用型
6.2 不通用
/ `. s* c* ]) f7 结语
1 图像分割和实例分割
/ O3 v; A( y% _4 D/ h& g& j( M  t7 a
( _* N4 H  I# U7 |* S* x& u计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
$ g+ V' r( B" g* S2 F4 j
" I* ^7 U5 v, F8 h( G7 n5 R( _# F6 L图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
/ }2 ~9 O% G  _" z) R" W实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

" ?- c2 Z* T9 ~2 Pascal VOC2012语义分割数据集

2.1 导入模块
import time
0 t/ u7 c+ l- b& I  _2 G: e6 Y& K- Nimport copy
import torch
from torch import optim, nn
; Q) \2 e* t' u1 d/ j. r5 c# I0 Gimport torch.nn.functional as F
import torchvision
" I; |/ d# I1 @$ E) d) Yfrom torchvision import transforms
2 J5 m, b% K8 E: E# C- l7 pfrom torchvision.models import resnet18
$ Q* W. C) t( ]* Eimport numpy as np
from matplotlib import pyplot as plt
from PIL import Image
  o, p" F4 v6 }! o+ _7 Uimport sys
) H8 _+ `8 t2 j; c2 V/ \. M* b1 asys.path.append("..")
& f3 j$ |0 N* x5 B1 g1 \from IPython import display
( \+ w6 o) f/ `6 x# _, A' ^8 }from tqdm import tqdm
7 g) j. m  j* l2 A3 Y, Eimport warnings
warnings.filterwarnings("ignore")
8 U, R0 ]! ?7 v
0 u+ E9 a  x; v7 M2.2 下载数据集

0 z, n- ]+ N0 _: n# q* S: K语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：


* [  k- B' H' E) [1 EImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
2 W3 n5 V$ N9 x) a4 o; |# MJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据
  E* C) B3 U' G
( a, U4 @, }% Y: F9 c7 B, l; A$ K0 i定义read_voc_images函数将输入图像和标签读进内存。
' Q/ H& l& E0 I
! i- ^; O/ V, y, ^; L7 D4 `' Tdef read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
2 o( Z5 d3 q) q$ X, D, F" i: n    with open(txt_fname, 'r') as f:
, X& e3 H$ ^, J7 I8 g        images = f.read().split() # 拆分成一个个名字组成list
7 X- C3 i8 L! o  l: N/ C2 o. W    if max_num is not None:
        images = images[:min(max_num, len(images))]
+ B. k. z8 J3 ]8 A' U. V7 q: W8 F    features, labels = [None] * len(images), [None] * len(images)
. x6 v  N  h  w0 s- g+ {: ^7 Y    for i, fname in tqdm(enumerate(images)):
) [* {- ^0 Q3 L5 R  F# y* Q* t7 \- U        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255
! ?- r4 w/ V$ W+ L" Q
定义可视化数据集的函数show_images
8 D' y; H! w! i* n- ]  f
4 N% q2 r  V  i5 g# 这个函数可以不需要
4 R9 i$ B' \8 ?! \def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
- t: j7 s# j" p: r& T# X    display.set_matplotlib_formats('svg')
: a+ J6 o3 E' c1 f  |    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
* J! T! E" n6 D3 a, d( ]7 A: e; F3 Z; v    figsize = (num_cols * scale, num_rows * scale)
( `1 g- m' G' `+ Q    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
. z* b0 ?9 e: u' n, v/ u        for j in range(num_cols):
( {# D- l7 V3 j/ w            axes[j].imshow(imgs[i * num_cols + j])
0 A$ j# c& s3 p; ^. }            axes[j].axes.get_xaxis().set_visible(False)
" B! x# g3 M# }0 P7 `4 W1 {            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
: {$ Y: j5 n# x
- s) l2 x, H# B$ f6 E; P- e定义可视化数据集的函数show_images。
5 P3 j5 X6 i2 M- H4 m
- w6 z) n# ?# M$ j, Q5 u: ?3 `2 C7 b- H# 这个函数可以不需要
& Z8 Y; _, d# A9 O5 Zdef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
2 U3 z/ g. G# y& R7 a    display.set_matplotlib_formats('svg')
: E8 Q3 g4 U* o    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
: c( W7 |& U( k3 k0 B$ K
! }; ^  F3 V% d* bdef show_images(imgs, num_rows, num_cols, scale=2):
( }- k5 Q* |4 s! I/ e- N% K3 y    # a_img = np.asarray(imgs)
% w  b& }; F  [" G    figsize = (num_cols * scale, num_rows * scale)
* q" j% `# Z& C) X    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
) h( v2 u9 \; [        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
/ i: `) H9 T' \8 ~8 Y            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
4 D( }; X0 x4 P1 P% Q    plt.show()
    return axes
3 p; M. H* t! W2 g画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# 根据自己存放数据集的路径修改voc_dir
* ^/ G3 h/ v" H" a5 k, xvoc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
2 ^0 r) y7 l+ e7 ?# ximgs = train_features[0:n] + train_labels[0:n] # PIL image
' ]' m: V. X' a. e. nshow_images(imgs, 2, n)


5 X7 S% F9 W% [& L
列出标签中每个RGB颜色的值及其标注的类别。
  P# ~5 w  D# C' T% r5 E# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
9 J2 j% p0 }2 k. `                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
* L0 T7 N# \3 x0 R& w                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
% m% Y' R' ~( b3 v- d2 l1 T8 L. i- x                [0, 64, 128]]
) J9 S, Z6 M3 U- |9 z# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
# H' X; b+ x" M8 o               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
$ ^! \* ^$ B. {  S' j, Q
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
0 w+ E* B! A. \4 D( j( H1 Ecolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
for i, colormap in enumerate(VOC_COLORMAP):
$ i; j8 f% h4 ^& j( x+ T    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

# 构造标签矩阵
4 H, V) |) d6 @3 Zdef voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
, M( B: V: L9 U- X    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
3 Q' |; M# Y$ k) H7 b2 w9 ]    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
( M. P4 T( n' l) N
9 R/ R5 X- N! ?; _) U可以打印一下结果

" v! L( c1 \* e# p9 e3 l- c' M# Cy = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
& A+ g/ \# R7 e0 |" S
2.4 预处理数据

在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。


def voc_rand_crop(feature, label, height, width):
3 M0 v& O) w) Z; D( Q& A) `    """
    随机裁剪feature(PIL image) 和 label(PIL image).
7 ?6 k0 m" {: s    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
        img (PIL Image): Image to be cropped.
5 O& M% A! t. e6 Q) l        output_size (tuple): Expected output size of the crop.
    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
  }) c$ d, l: l: y9 J# J/ q3 ~0 Z' H9 M    """
0 ]/ f2 @( x* m6 m5 r8 a    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
6 |; \6 `. |1 r    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
% m2 h/ G  z8 Q+ Q    label = torchvision.transforms.functional.crop(label, i, j, h, w)
$ ]8 @+ c, s+ W0 N5 e    return feature, label

" B; k4 Q. b! y& i, l0 N# 显示n张随机裁剪的图像和标签，前面的n是5
+ t, h0 n: Q3 v/ H5 m( Nimgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
0 g# T7 e6 p8 m! ~/ }$ j. R! cshow_images(imgs[::2] + imgs[1::2], 2, n);


& O! ?* u4 c1 @
( Q6 H1 ^3 P9 r

; f% }& V2 k( h
3 自定义数据集类

3.1 数据集类

torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
# g' I% h: |0 H2 ?5 \) g
__len__ 实现 len(dataset) 返还数据集的尺寸。
# l' W  _7 l) z__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
2 U% Z! L9 f: H; B2 k( `由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

4 M9 v* }# l* u) e  v7 Oclass VOCSegDataset(torch.utils.data.Dataset):
) p4 N5 |4 q, S! v    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
# v2 s2 J$ N# b3 ^& J        self.rgb_mean = np.array([0.485, 0.456, 0.406])
+ z; ?4 F' V  \, d        self.rgb_std = np.array([0.229, 0.224, 0.225])
        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
# K) ~$ E/ H7 R( K- I* a8 I            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
6 F( K6 h2 C2 }/ R' e        self.crop_size = crop_size # (h, w)
- B$ ?. [; ^' F" [& V* `# s        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
9 K* Q' ]7 d/ E, o3 b) W7 ^        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
. n2 N+ J$ N6 R  ~3 g; g3 j* Z        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
3 ?6 c9 M9 F0 z# P+ S" M
    def filter(self, imgs):
" S2 g& A7 H) R) U  X        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
8 O' i0 {% s. ]8 S* C! \9 ]
    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
                                # float32 tensor           uint8 tensor (b,h,w)
* B6 w! ^: r& C2 p; j5 r3 ^" @3 T+ o, G        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
. j, S7 K/ I, i
! B" b7 E& ?. K, L! k4 y4 F" s    def __len__(self):
        return len(self.features)
: |, O: z5 x. F3.2 读取数据集
, u& B8 T% \! g+ p( V
& e% ]* d1 M! U  o通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

- P  t3 {3 X! O. V$ t) E5 m! I& vbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
# X9 j. ]$ _3 i* Z- N( M/ t8 Ymax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

2 W4 |3 z" `0 s6 ~# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
: L) ~1 a5 b& z3 S" Atrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
2 s# F5 K; _$ M. |+ ]: X, F/ w  ltest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
$ R3 Q8 ~- p. {6 d$ ^: p" d. t, p( t                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

6 A$ Z# l+ _; J$ q4 构造模型4.1 预训练模型

下⾯我们使⽤⼀个基于 ImageNet 数据集预训练的 ResNet-18 模型来抽取图像特征。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
! k" F8 }; u" Z, c6 _, b+ G
num_classes = 21 # 21分类，1个背景，20个物体
4 X2 y4 P' r; Q* l: c7 d0 Lmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
' j3 r0 `& Y0 _, h" ]0 _$ ]7 q) Z
/ m: b9 X4 b9 z, D# 特征提取器
2 D$ y% m4 d+ L0 B6 m' bfor param in model_ft.parameters():
9 P/ d- A" R6 D4 [% L    param.requires_grad = False
$ d/ v3 I) r. V: ]8 y, w! D) r: {: P4.2 修改成FCN
  b3 p5 Z5 }: ~- g' e( t
全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
3 }$ d' c$ b8 s8 }+ F最终输出的通道包含了该空间位置像素的类别预测。
/ W, I$ P5 s" q& T4 ?8 n$ M
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
4 \& W# x! N4 |+ \
可以先打印model_ft，可见 ResNet-18 的最后两层分别是全局最⼤池化层GlobalAvgPool2D 和 全连接层。全卷积⽹络不需要使⽤这些层。通过测试，当输入图像的 size 是(batch,3,320,480) (batch,3,320,480)(batch,3,320,480) 时，通过除最后两层的预训练网络后输出的大小是 (batch,512,10,15) (batch,512,10,15)(batch,512,10,15)，也就是 feature featurefeature 的宽高比输入缩小了 32 3232 倍，只需要用转置卷积层将其放大 32 3232 倍即可。
: Q' \) e" c# d
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
; V: Q' V" P" d( s% ~' l              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
( T1 C4 Q$ C/ I6 W, P
( x+ A" V$ d4 z" l# _" _: U# 对model_ft做一个测试
4 z% I( ?% {0 t. mx = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

8 Z$ p5 `5 m0 F" ?9 ], h# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
#     break
0 u+ ]6 `* c) B: x" G0 d

$ l& S$ k5 n$ ~& @/ d5 O* K2 u% D4.3 初始化转置卷积层

/ n: a9 r7 r( N4 x在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
6 W/ B4 `4 K+ D* x6 p6 }+ N/ u/ s4 G在坐标 (x,y) (x, y)(x,y)上的像素，先将该坐标映射到输⼊图像的坐标 (x',y') (x′, y′ )(x′,y′)。例如，根据输⼊与输出的尺⼨之⽐来映射。映射后的 x' x′x′ 和 y' y′y′ 通常是实数。然后，在输⼊图像上找到与坐标 (x',y') (x′, y′ )(x′,y′)最近的 4 44 个像素。最后，输出图像在坐标 (x,y) (x, y)(x,y)上的像素依据输⼊图像上这4 44个像素及其与 (x',y') (x′, y′ )(x′,y′)的相对距离来计算。双线性插值的上采样可以通过由以下bilinear_kernel函数构造的卷积核的转置卷积层来实现。


# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
        center = factor-1
) m' v& [9 i: G% c6 Q6 J    else:
' }$ @; m! M# p# n        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
3 X- x' l7 W; q5 T; }3 Z5 P2 Q    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
1 ~: E3 |/ z5 q$ I: R& f+ A    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight

& D$ i( k/ M6 R7 G7 V
7 E/ `: m6 `2 |2 ^# G在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
) Z# O# R8 {5 t) j2 S- v
nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
: d6 g4 }" D, k. i" N/ j. B
* j, t" _3 G+ H! G/ o8 J3 N

( u2 _# d% }7 g2 R( w, W5 训练模型

现在可以开始训练模型了。这⾥的损失函数和准确率计算与图像分类中的并没有本质上的不同。有一个 blog 我认为说的很详细，图也画得很好：https://blog.csdn.net/Fcc_bd_stars/article/details/105158215

0 ?" A7 K  Q1 Z% L8 [def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
' ?6 @! K/ o8 C  O9 ]: k    # 每个epoch都有一个训练和验证阶段
3 P" O% k3 F. w% J8 N  _    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
        for phase in ['train', 'val']:
+ u" p3 f+ Y5 a5 p7 c1 @6 u; L! X            if phase == 'train':
                scheduler.step()
6 C: ?3 Z3 E. L( G                model.train()
- p. |4 L, w. @; ]6 w# v            else:
* r! {' I( o( w0 M% }                model.eval()
$ b/ L2 ]2 I4 E# |, }3 _            runing_loss = 0.0
            runing_corrects = 0.0
            # 迭代一个epoch
: s+ t* L7 g9 ?; z            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
1 b* o( i/ \* W7 A, o                optimizer.zero_grad() # 零参数梯度
                                # 前向，只在训练时跟踪参数
6 _( C- ^3 `2 A' m" `  ]1 L8 n- L6 |* W( i                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
% u0 V+ ]: Q, F                    if phase=='train':
% ^7 f# j  ?, Q5 B# [                        loss.backward()
+ b  h0 E7 b+ r# J( \                        optimizer.step()
                                # 统计loss和correct
/ `$ n3 w( h# M- x% k. o, p                runing_loss += loss.item()*inputs.size(0)
2 x' R& D. o, r) Z& R3 j( l, [0 Q; A                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
2 t! z0 p# `  P: m4 d5 f# c8 ^
            epoch_loss = runing_loss / dataset_sizes[phase]
% `; M; y& b+ M+ ~. S            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
" A- v9 D5 j4 ~. L7 W            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
4 p7 [7 C% W" @% }, |                        # 深度复制model参数
) K1 U, _* R0 y1 O) N4 t: \& g            if phase=='val' and epoch_acc>best_acc:
6 |$ Y2 J( b- E7 e) w" n                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
" j) V7 q, L" O: ^: U- p        print()
& F* Y7 Z5 c, k$ B2 v    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    # 加载最佳模型权重
: n$ }: n, s% O" ?: @    model.load_state_dict(best_model_wts)
    return model
. K: |8 `0 W" b8 q
" v3 e3 E( S& N$ |7 b3 T$ ]$ g5 y$ T下面定义train_model要用到的参数，开始训练

epochs = 5 # 训练5个epoch
0 E+ n2 C0 `+ _; V* w8 pcriteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
1 S" [. C3 P$ ~# \0 nexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
6 d; o* N* f9 e- G7 G2 T9 L
4 Q7 c" C+ p0 G  S1 O# 开始训练
: j, S5 w' S3 B) w7 _model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
0 d8 G  T$ d" H
6 测试模型

为了可视化每个像素的预测类别，我们将预测类别映射回它们在数据集中的标注颜⾊。

def label2image(pred):
    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
0 k- i# b0 E  W# Q5 K5 u    return (colormap[x,:]).data.cpu().numpy()

下面这里提供了两种测试形式

6.1 通用型

其实如果要用于测试其它数据集，也是要改动一下的 : ) 😃

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
7 y% f5 w4 o8 U2 c2 j5 \std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
6 i+ a5 ~) h3 kdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
. `" a( h+ C4 i0 |' _7 m1 t1 J    model.eval()
    images_so_far = 0
  S! ]2 d* p# K* F    n, imgs = num_images, []
* Q1 m0 E+ x  r$ L' c    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
  w; }* n' J+ d3 N6 `& w6 y            outputs = model(inputs)
            pred = torch.argmax(outputs, dim=1) # [b,320,480]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
6 E+ R" V' ?: j+ X. V+ o: f& C
            for j in range(num_images):
                images_so_far += 1
                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
. w+ i  e5 K5 H3 j) }5 G                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
2 A0 P5 l) ~9 T% q# N  N& d                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
$ j# |) F( |7 T                    return model.train(mode=was_training)
; b3 p( t3 v1 N% g
% H5 D8 t5 H$ p8 [% V# t# 开始验证
visualize_model(model_ft)
6.2 不通用

在预测时，我们需要将输⼊图像在各个通道做标准化，并转成卷积神经⽹络所需要的四维输⼊格式。

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
0 i* r; [7 d& _2 P2 C            transforms.ToTensor(), # 好像会自动转换channel
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
" [0 a! U3 N. X3 e9 \' s8 K0 @
: Y3 t$ o, |- _; f, _5 a9 |def evaluate(model:nn.Module):
% v* N5 x! f0 |/ H    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
' M: _+ x5 I8 y5 P( K    n, imgs = 4, []
" M0 T3 o% B1 E! M2 F, n$ m7 J    for i in range(n):
        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
2 O: o' s$ I- m        pred = label2image(predict(xi, model))
* x, ?2 N+ q+ o# i: v; X$ H; }# j        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
% j9 N* O5 a& c# X- x1 Z( n+ a( {
# 开始测试
evaluate(model_ft)

1 b% v+ k$ e3 C, R% c7 结语

我只训练了3个epoch，下面是训练输出

Epoch 0/2
6 ^# ^* a3 K" L0 }----------
/ K8 l% I' H9 U: J9 t: otrain Loss: 1.7844 Acc: 0.5835
$ O* a2 Z0 {) U; r3 H( g8 k7 K* `) \& tval Loss: 1.1669 Acc: 0.6456

' i) C) V5 A* q4 }1 dEpoch 1/2
----------
( u; y& @7 W( Y. f. \% mtrain Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

Epoch 2/2
----------
6 q+ K  ^# Q' j3 ]# wtrain Loss: 0.9578 Acc: 0.6706
9 o( {' A8 I6 M" A. _& c* Zval Loss: 0.8088 Acc: 0.6948
9 A1 C, l! h; K4 s
Training complete in 6m 37s
- a& L0 w' u4 n: ^7 z$ E5 z

. g- _' f9 Y/ u0 {3 d0 S+ ~" c
2 F" X+ x6 E% B+ s) i* T% l9 s9 N
& m" w. Y6 |2 G) A% ?: U9 `当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。


8 y7 O  S  v- X0 R语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
/ p9 W% K8 D0 ~2 GENet 用于实时语义分割的深度神经网络体系结构[论文]
9 `5 E4 J% p9 _( G: i/ U* e% b+ kU-Net 用于生物医学图像分割的卷积网络
( O- D. M8 T$ OSegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

; m( d3 ]! S1 q6 }) s- u常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
2 p* \+ {8 o5 U4 x7 V
对于损失函数，除了交叉熵误差，也可以用这些：
/ L% ]6 \  N+ R9 O
Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
5 c: t. q" g6 p! e" _
; A7 D9 l9 I9 N5 E- c

: g, u: U! w# u. {% [
1 j6 O* U" G( V/ {7 D
————————————————
$ K/ |( |! c1 v5 C. P0 d+ h版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507