Pytorch实战语义分割（VOC2012）

杨利霞

1 C# Q1 @& G9 N1 {4 J9 ]Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
$ x- @1 G9 V' v& c; a语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
  b5 C7 `/ A6 C# H
' e! ~0 J) b7 \) _" K7 \3 Y
0 J1 t- s3 I- k2 F语义分割中图像有关狗、猫和背景的标签
" Z9 o1 s# f- }4 e文章目录
2 u5 p, G% @  U2 J( T+ U! Y
- q4 L/ r+ D, U+ o4 k1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
2.1 导入模块
' a" U, z0 S9 v: D- X2.2 下载数据集
2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
( R: Q9 m8 `2 j8 F4 F3.1 数据集类
- t" a- M4 w/ y# |3.2 读取数据集
' c) ~9 g5 y  r( m, K8 t! h3 `4 构造模型
4.1 预训练模型
4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
6 测试模型
7 X. [! Z7 w" N/ t" k# H& J1 A6.1 通用型
6.2 不通用
7 结语
; ^5 v; U5 ?8 N- p  f, D( y1 图像分割和实例分割

3 l! o7 N. `' [6 Z计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
$ E) O" _5 v! B
图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

2 Pascal VOC2012语义分割数据集

2.1 导入模块
! \2 C' X  F4 h( O4 Z9 {" ~( eimport time
import copy
import torch
from torch import optim, nn
5 o9 `; }0 n4 w% Z* r( Bimport torch.nn.functional as F
/ F5 B9 Q7 k* K* d4 \$ M3 ximport torchvision
! s. t4 }7 t8 {; O( o9 Nfrom torchvision import transforms
; e" I7 ^- L7 ?+ y& H4 L, Vfrom torchvision.models import resnet18
import numpy as np
from matplotlib import pyplot as plt
0 F' T0 @! O) ~& |2 X, p' X2 hfrom PIL import Image
import sys
3 t' ]1 g# W) H6 I4 s4 ?sys.path.append("..")
from IPython import display
2 w1 z/ E4 P' I9 i6 l4 \* wfrom tqdm import tqdm
import warnings
+ b6 G, r0 H' wwarnings.filterwarnings("ignore")
2 m! D9 I$ F2 B1 O0 u- d( C9 T
; K) p8 v4 @* r* l7 c3 t2.2 下载数据集
; C0 s7 z& X0 w' p$ P- q0 ]
3 H3 u0 b; e' z% M, l语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
% Y  ?) D  [. |

ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据

$ O8 W+ p) W3 r9 r  i定义read_voc_images函数将输入图像和标签读进内存。
9 H7 w% K9 F% q) [6 J
def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
, y+ f$ {) Z6 M  b6 N4 b* _, t9 }  v    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
    with open(txt_fname, 'r') as f:
        images = f.read().split() # 拆分成一个个名字组成list
- b$ s. V/ M) s$ f! J    if max_num is not None:
0 z6 V* ]& k+ r5 P  E        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
/ B5 a* Z; B% W3 j6 s* b. u        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
) a* G9 P) n* K$ F  n        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255
3 E" m$ A; w9 f( `7 Q0 G7 ]
定义可视化数据集的函数show_images

# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
9 {5 p) L9 {4 [! P5 _+ R' U    """在jupyter使用svg显示"""
3 V2 }0 V( X6 P) {2 Y  |    display.set_matplotlib_formats('svg')
. e$ \3 V4 A' _( @) J0 F  S. x    # 设置图的尺寸
) r. |. i! q, O% f! ^# Z    plt.rcParams['figure.figsize'] = figsize
( G" j$ I. B5 C
% e8 a6 x, O0 C' e6 n  wdef show_images(imgs, num_rows, num_cols, scale=2):
3 B2 A. H. ~5 \4 m    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
! n- `! j. L8 g, G/ e    for i in range(num_rows):
" e: ]; {5 l5 s. j# C        for j in range(num_cols):
7 [2 @# V1 J3 F4 b$ s4 B% G9 r) r$ r            axes[j].imshow(imgs[i * num_cols + j])
3 M% {7 d- N: M- ]& q4 H  z            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
4 [; p6 s& L0 L: D$ T    plt.show()
    return axes
7 A3 Z9 k5 x; S: m
9 U, L/ s2 t% c2 E/ e2 F定义可视化数据集的函数show_images。

# 这个函数可以不需要
& {7 Q! h  b! C% M6 ]$ xdef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
1 [: G: A9 i" F( ]7 |7 D    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
. V  m, F7 K0 I
def show_images(imgs, num_rows, num_cols, scale=2):
1 ^* h; Z' A' {9 a" G    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
# M! u7 J8 T' a    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
7 b6 |! }3 T- b4 `6 H- D: Y        for j in range(num_cols):
& C! o, E' K. l' ?7 e- P* ^$ g            axes[j].imshow(imgs[i * num_cols + j])
6 I5 \  C8 |6 O$ e; v" q( {            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
& d  R5 ~& w' ]* c$ j2 p    plt.show()
    return axes
0 N( H/ L/ _4 l9 U3 o画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
& u/ t* R+ e: u; l
8 f2 b* s' g& A; b8 ]: |# 根据自己存放数据集的路径修改voc_dir
8 G. @9 f9 j  hvoc_dir = r"[local]\VOCdevkit\VOC2012"
. o1 r# G9 `+ C) u, p# o' J1 ttrain_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)

3 Z) V0 \% D) F# T; d
1 H% X  M: l9 @+ b2 O
列出标签中每个RGB颜色的值及其标注的类别。
& D# K, W% ~6 A" F( g) L5 x# 标签中每个RGB颜色的值
VOC_COLORMAP = [[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],
& L# a% I- w2 P% X* s; Y1 D                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
# V6 I1 }+ v) g2 v                [0, 64, 128]]
: `$ k. }  G; X* p. o# 标签其标注的类别
; c9 M- R) P: I1 Y5 J/ l9 KVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
# y0 T5 \+ `0 W- L5 z1 u               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
# O5 ?% o" o9 I, w4 S. Q1 X, e               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
) [, h, p* i4 g$ R
4 d( r7 w& X8 ~* _$ S有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
& o9 p* p) J4 u8 v- O' x) Vcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
* C7 X8 O* Z) w% P+ ~for i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
& y, @4 O4 v. O1 k0 T3 T& P    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
5 q$ s2 w  l" V& B3 P4 b
# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
' |! w" w3 K& R    colormap = np.array(colormap.convert("RGB")).astype('int32')
3 K& k' R# T% f( J) i    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

可以打印一下结果
' e* N; e, g( \& m! u$ A
y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据
5 a6 ]6 |& N# V" g$ x
( `1 A1 f1 C- {! }在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。

! c3 y7 I7 g6 a6 P2 {* J0 Z' d
def voc_rand_crop(feature, label, height, width):
    """
1 v" A8 q! `) `    随机裁剪feature(PIL image) 和 label(PIL image).
: Q4 Z2 `2 q7 C    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
8 M" b7 L% l5 }1 _    Get parameters for ``crop`` for a random crop.
    Args:
  O( {; N& J+ W" t        img (PIL Image): Image to be cropped.
/ H9 R! Q( o; L: F! D/ Z( C        output_size (tuple): Expected output size of the crop.
    Returns:
* N7 }4 T) ?9 ~" X$ P- N        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
* |0 i" O% Z/ v  E, |- H    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
+ |7 |5 C, A' s9 w+ t4 T    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
" n' V! u3 J' W    label = torchvision.transforms.functional.crop(label, i, j, h, w)
5 p3 A% q: v: f! o1 J    return feature, label

' m  h+ c4 b/ B: K; I# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);
3 u+ ^( ^2 k  r6 m8 W7 N# L/ ~- Z
+ k; X- k$ o1 M. h' l/ y

" W, d: U' i6 Y0 F$ c" ?0 t
! ^. _* ?0 g0 R5 ~

# h/ Q* T, x( X3 m8 E. e5 J, M3 自定义数据集类

3.1 数据集类

1 \5 W0 Z0 ]9 {! b* ~# r$ N* ptorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
9 e0 q" C6 O, e3 _" ?, C
* b; Z6 i7 U! p; Y" D. u__len__ 实现 len(dataset) 返还数据集的尺寸。
6 I" G8 h% e+ l$ T, X6 z: a__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
5 J" d$ @- C0 k) k/ H) Y由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

' m" [/ F! ^5 Bclass VOCSegDataset(torch.utils.data.Dataset):
! x: g+ i5 l7 }& A- w9 [, f    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
) S5 V; ]  x! P+ B: r        """
        crop_size: (h, w)
: P- ]* d9 D# `( Q" e! j        """
* L' E$ t4 o  u2 O0 k        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
6 Y( C. v' W. c: U- E6 N, |! u8 b, H        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
9 _4 K5 g1 A0 I# y" m: ]# F# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
  g4 _/ D( M5 u4 Z9 ^4 L6 \        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
0 J6 w9 z& E/ o1 X% M2 W
9 C1 t. j; F+ f) \. y% C2 N) U. j    def filter(self, imgs):
        return [img for img in imgs if (
' ~$ r; x# E/ }5 I5 X- {            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
; f: v- P1 x! n+ w
5 r. M% ^5 I: _2 |# _9 r    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)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
: H/ R, p, t5 t1 X* K3 f/ f( C: z" v
    def __len__(self):
        return len(self.features)
/ }4 o  }+ R$ z3 \8 N6 U# u3.2 读取数据集

8 Y! ]* `! Q# u通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

8 {# h- |6 r7 W$ h+ v6 xbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
3 T' l& P- N" e: q5 _8 [
7 e' _  n$ g4 ]7 Z! E, e; j( r* I9 D# 创建训练集和测试集的实例
9 z& t8 q6 T1 m% avoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

* ~3 o2 v* b+ D; r# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
8 W) K! ?4 A) G                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
# j( U' V) F* H4 M' V! D- L* f                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
7 M) }# X- o% i# bdataloaders = {'train':train_iter, 'val':test_iter}
8 A2 C3 y$ ]& Z2 n/ e7 v' m: ]# adataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

num_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

" e4 q, T2 s0 O0 {/ x# 特征提取器
+ k. b$ A/ m' S  n0 ?; ?for param in model_ft.parameters():
    param.requires_grad = False
! X5 V6 ?1 [' |; A- b4 g4.2 修改成FCN

全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
, p3 O! U! W8 |  s+ R最终输出的通道包含了该空间位置像素的类别预测。
( A( m# j9 |: `1 g/ I$ z
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
* I& Y% j1 z1 m: ]  N9 a2 S9 _
4 Z5 D. u  z. K0 p' o可以先打印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 倍即可。

7 F! p  E2 t& I; e; X+ H! _# U7 ~model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
* [3 c0 k6 K% h- ~  s0 Y              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) # 转置卷积层使图像变为输入图像的大小

0 F" x& i% J0 a! V/ K* M' Z, R  H# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
) I0 w/ k. F5 h. Lprint(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
: j# ^2 ?9 @. V9 j
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
' C9 f: M0 g, \" \# for X, Y in train_iter:
  A* s0 D& s) U#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
#     break
$ e: M/ O+ \+ U

4.3 初始化转置卷积层

- R! b4 T# {* I在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
在坐标 (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函数构造的卷积核的转置卷积层来实现。
& T: d' a, h9 b, o7 B8 p+ g
% Q- ]- T. b8 v% S
: U5 s# H/ u3 A' `) L( k# p* y# 双线性插值的上采样，用来初始化转置卷积层的卷积核
/ O  U* U+ b" Ddef bilinear_kernel(in_channels, out_channels, kernel_size):
  w9 s0 H& m0 K3 L    factor = (kernel_size+1)//2
3 g8 K4 {$ \+ C& u/ Z& V    if kernel_size%2 == 1:
        center = factor-1
    else:
        center = factor-0.5
3 k' B9 @, j  S4 X    og = np.ogrid[:kernel_size, :kernel_size]
3 t5 r& ]0 m/ b, c' C, E7 s& J    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
, k: o7 b8 A8 U+ M. G4 p. F    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight
# T7 _1 Q( m& D( ]4 _

' S$ M% k# t, x7 n, \在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

8 z9 b  w+ P, w. \& Xnn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
, P8 D' B% H0 z8 }- Z/ }9 M8 cmodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
! |  x6 Z" t7 f) Q0 |* Z7 k
, h  I) s- s9 j  y- L9 ^' z

" W* y2 M4 E% f" A$ s% R  y) u/ ^5 训练模型

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

2 T& F- O0 B- q7 i  b' ?5 s* Tdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
( `4 O# ]% h) w0 X    since = time.time()
9 U' R! a& Y3 a8 g/ X( S    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
$ x% X5 g4 w' ~# q' u    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
; ]: R% r$ L. N$ P0 [        print('-'*10)
3 m* k8 f# v; U0 K7 {. x9 U6 F  b        for phase in ['train', 'val']:
            if phase == 'train':
: M& o" G1 c% u& `. F: z                scheduler.step()
0 Y- v9 x" t$ `                model.train()
" F/ v$ @  T, Z9 g            else:
6 ~' H, Q4 {, u& I                model.eval()
' L4 I4 K1 N# B/ U            runing_loss = 0.0
/ ]. M) F, u) W1 a: o, h3 t" O            runing_corrects = 0.0
            # 迭代一个epoch
0 g* u' T( T' N; K2 Y+ q4 ]$ `( e            for inputs, labels in dataloaders[phase]:
9 y9 w1 u* I% k, T: L3 x                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
" t3 A4 d# B1 ]/ Y8 l" N- F                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
0 u& R' x2 H4 B; m* ~* g                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
                    if phase=='train':
8 y6 C5 |4 e% {% k# e2 t                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
% _# Z2 N& S2 T/ d1 A                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
: U# Q3 t3 O* Q0 u
            epoch_loss = runing_loss / dataset_sizes[phase]
% J) ?3 c& d' y+ c: ?& _8 W            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
. v1 @5 k9 {- k  K8 S                        # 深度复制model参数
( `  V! e) E/ h0 v3 K. `            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
, h9 y' ~. X3 B                best_model_wts = copy.deepcopy(model.state_dict())
; C3 Z3 c% h2 h  T" Y2 }        print()
- j7 ~  T% D5 h1 Q+ W) }! I    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
/ a& w0 ?- y6 Y3 d, a4 [    # 加载最佳模型权重
! _2 k+ Q' q0 s    model.load_state_dict(best_model_wts)
    return model

5 e4 m, ?5 w6 G$ J) h; M' G) N/ q下面定义train_model要用到的参数，开始训练
5 T; ~. |6 u: a4 g& M5 k/ ^
epochs = 5 # 训练5个epoch
  C  Z4 t3 d1 _9 s' Kcriteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
: v" u) E; _. k0 g, N; u# 每3个epochs衰减LR通过设置gamma=0.1
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 开始训练
9 G* y. i/ J! E% f! j+ Rmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

6 测试模型

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

def label2image(pred):
2 L9 q2 l8 m+ r, Y3 [! |. l0 C    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
4 ]+ ^0 ]* Y2 H2 I& ?8 \$ ^2 H1 L: U5 G    x = pred.long()
3 p/ B& \2 ?6 r/ A' m    return (colormap[x,:]).data.cpu().numpy()

1 a3 H& L+ J0 x2 i6 r: N

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
7 Y. M' L" k" X( ]- c1 [* p' Pstd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
def visualize_model(model:nn.Module, num_images=4):
" W* ]0 {1 ~2 _! o( R# F( [% r    was_training = model.training
1 M5 j+ T( K' M6 j0 X    model.eval()
    images_so_far = 0
2 M7 Z1 V5 i" i$ ]! G; c/ t& p    n, imgs = num_images, []
, t: \% \. h( g5 C0 s    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
4 ~) b* V- E1 ^" n            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
. W# C; C# h2 q% s' m. k4 t$ `  N5 T            outputs = model(inputs)
( K) r& |4 D) E7 D9 P/ \            pred = torch.argmax(outputs, dim=1) # [b,320,480]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
& f; U7 d' H( t, n6 J6 ]. A$ u* s; @
            for j in range(num_images):
                images_so_far += 1
3 ]2 Z, ?$ g9 }9 h% w) }  Q                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
6 e8 C/ D2 M: }                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
% X" G% E5 F- j9 V1 s/ }$ m3 M                if images_so_far == num_images:
+ s8 o7 i& o# p# |                    model.train(mode=was_training)
+ d# @7 L/ E5 a" B6 m7 b4 o                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
; b4 |. L5 f, W3 i% l; m                    return model.train(mode=was_training)
9 m& Z6 c% d+ V5 V
# 开始验证
visualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
$ Y) p" g3 h# v1 N! A# {7 \6 R    tsf = transforms.Compose([
+ h/ |! T) K/ v/ P            transforms.ToTensor(), # 好像会自动转换channel
% Q% J" R& x9 R% `2 {            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+ R7 M) \4 b/ h& d- E; X    x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
' h- b& T) Q2 X/ I8 Q% w/ A! ]    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
3 y1 F" m0 m7 D* X
def evaluate(model:nn.Module):
: {6 S* s( x. V- H  ?6 [* ]    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
    n, imgs = 4, []
: F% L" I) T9 K6 t    for i in range(n):
' Y4 ^+ u0 _- c! C) |* }' d        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
1 h% \% m: }) F; e; Q- e0 m        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

# 开始测试
# r& Y! K( w5 n8 Bevaluate(model_ft)
3 T, m8 y( m/ J0 j* F0 |9 i3 s
% @; B+ q0 J, I6 R+ S& U, i7 结语

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

Epoch 0/2
----------
8 a  b! b- v7 Z0 C1 F; F8 ^8 Rtrain Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456

2 h  q' ~+ d" N5 HEpoch 1/2
----------
8 j* i1 ]9 j9 D* O* y: Ztrain Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

Epoch 2/2
----------
train Loss: 0.9578 Acc: 0.6706
% }& e8 x7 _- G( x$ r3 I& R4 Gval Loss: 0.8088 Acc: 0.6948
8 W$ C, g! N, ?, R$ u
Training complete in 6m 37s
( G0 q& C7 |6 S6 @) f. P
6 M$ i; s' P+ J* V* Z


当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
% g6 s9 p& d9 ]  b# z) f, P( e
. v; a% N, }  ]" G
* T3 x# Z* n9 V- q语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
# Q3 R0 _3 M) q* W# ?9 e
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
- ]! j9 F$ w% T  L5 o. mUperNet 统一感知解析
. x3 N) \; B7 f* vENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
( s2 C2 G$ M0 A0 i' Q* w7 N还有（DUC，HDC）、PSPNet等。
, v+ ^1 r) A4 _
) V, t/ H6 }- N( [常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
+ E) M7 y' m8 k) `2 H$ d8 O# {
对于损失函数，除了交叉熵误差，也可以用这些：
% R+ l6 V, z. U5 c8 |# x# v% a
' w" A8 n7 g& D/ dDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
9 p- }$ A, L( B' _CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
( W  Z4 P, u9 xLovasz Softmax 查看论文：Lovasz - softmax损失。
; A3 J) s' X0 g9 E! N! c1 C  V
/ z) V' j# F8 K% G3 y* E2 m

. D- }4 Z0 {2 ~, b6 v8 Q
; t8 M4 a  K5 F' P
————————————————
! O8 f8 N: ?/ ^- h$ d# V+ ]( m. M版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
: G3 w; L1 t  O, B' w$ n原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507


4 A: M: v) F; x; ?. v