Pytorch实战语义分割（VOC2012）

杨利霞

5 h% W) G3 |: J4 G. m# P6 X
8 Q5 c. ^& s+ x$ NPytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
  I% a% ~' ?3 U. n5 ~语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。


5 f+ h5 k3 q- S5 X语义分割中图像有关狗、猫和背景的标签
& r8 \; n; I& t9 r1 H文章目录
3 u. R1 q6 h* e
1 图像分割和实例分割
. E" ^* J* r' X2 G6 Y- m2 Pascal VOC2012语义分割数据集
2.1 导入模块
- f9 S1 n2 b( b* B9 W2.2 下载数据集
3 C& V6 M' J2 L2.3 可视化数据
$ U* f8 A4 I* h9 _$ S2.4 预处理数据
; y7 u0 p; X: \$ r; W7 B3 自定义数据集类
. p6 A' N% l) V* F6 z3.1 数据集类
3.2 读取数据集
; l. ~% r2 z$ T5 ~5 K0 [4 构造模型
4.1 预训练模型
4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
6 测试模型
6.1 通用型
. J5 N0 {' {  G6.2 不通用
7 结语
1 图像分割和实例分割

3 }# c0 N# S$ R3 Y+ D5 p计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

# n5 i) ]6 V5 ]( U" V, R: n图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
) X9 E  t  t6 B& L% T; h1 A- b/ Z
2 Pascal VOC2012语义分割数据集

8 k# }* v, N9 S5 ~5 D2.1 导入模块
import time
! V+ f6 x  z: a4 `' oimport copy
9 y- ]$ P( m  W4 Z; }: ximport torch
0 j" g6 V6 t! ~! v" S1 L0 h: _6 ^7 cfrom torch import optim, nn
, g8 s3 @0 F* T, A. W! E; Qimport torch.nn.functional as F
2 T0 F( D* W# G9 kimport torchvision
from torchvision import transforms
from torchvision.models import resnet18
import numpy as np
from matplotlib import pyplot as plt
: n# ~5 I' R" Y. A6 P4 Dfrom PIL import Image
3 F: e2 m  |5 oimport sys
1 l$ n6 }) K$ Q# d% Q; G' [+ v; Isys.path.append("..")
5 r- v, u# L, D$ {  m9 ^  a8 zfrom IPython import display
from tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

2 y+ R: x( O% _2.2 下载数据集

语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
3 s2 g& z2 \$ U3 P

$ G8 E" y4 W+ o$ C# rImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
" S! H4 p( }  LJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据

定义read_voc_images函数将输入图像和标签读进内存。

3 I+ W, b3 F2 cdef read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
: \( ~, W; o4 E9 }$ @1 D* i7 L% V' x    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
# g( D6 j1 @5 S9 w    with open(txt_fname, 'r') as f:
3 ~4 B8 j8 R8 e$ }2 _9 r  @# ?        images = f.read().split() # 拆分成一个个名字组成list
# M& O0 x2 g* v  N. u    if max_num is not None:
        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
9 [( j$ w$ e, M! T3 D- k& O! ^    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
1 C9 b7 t- F4 r7 X% s        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
1 R' m  y( g3 J, m! b        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
8 F% f" c+ A+ ]: \+ i    return features, labels # PIL image 0-255
# R* j" ~9 C9 ^
定义可视化数据集的函数show_images

2 |) O! t. a. s4 W: e# 这个函数可以不需要
& Y' p. d+ e6 A2 tdef set_figsize(figsize=(3.5, 2.5)):
0 }2 p. G7 ~# n" u( o1 {9 s) W    """在jupyter使用svg显示"""
  W4 a* `0 g) p. c9 u- y    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
& [- P7 C- F% I8 y( l% K    plt.rcParams['figure.figsize'] = figsize

. D' ^( w7 _: a- Tdef show_images(imgs, num_rows, num_cols, scale=2):
# N; c$ Y/ d: M1 i    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
. y! s; ?  e. H- f/ B4 Z. b0 R% X5 F            axes[j].axes.get_yaxis().set_visible(False)
% E6 ~4 j8 N2 C6 r# R( R' e& T1 t    plt.show()
    return axes

6 C1 T) i4 W3 _8 l& l8 \定义可视化数据集的函数show_images。
/ W" ~5 h& f1 D. n% U, Q7 N6 q
2 ?) C3 ]* Z( |& L& W1 ]8 t/ @! j# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
, ^  T9 Z4 A  q    plt.rcParams['figure.figsize'] = figsize
, s! K) X" L5 v# f
2 U8 v5 }7 g# e6 o/ L' l7 ]) @def show_images(imgs, num_rows, num_cols, scale=2):
! A2 d! d# c! v6 X2 v% |7 T1 E    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
; ~8 l& E7 r" I" i0 s    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
9 z+ {. J' h- n  X% R    for i in range(num_rows):
9 z  w- T7 I. X4 B& \8 M( G        for j in range(num_cols):
' F% n5 L2 X# r* v4 U+ ?' Y            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
$ H$ Q* i: h/ s( p    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
8 J% h% \6 w6 I# r# T6 u
; ~+ h; N# g& N' A; ~0 I0 p- z5 H# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
0 Z2 O4 {$ B( `. a* htrain_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)
+ N, A# v. S8 K- F+ u

: ?3 {8 O# [+ y# Z0 n
. ^$ O/ L- C% e列出标签中每个RGB颜色的值及其标注的类别。
( z; ~$ p7 P7 H# q1 m# 标签中每个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],
                [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]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
! V1 D. y3 l+ `1 l+ t- G9 }) j2 M/ |( l               'diningtable', 'dog', 'horse', 'motorbike', 'person',
8 u/ s- h, D9 Z  W% Y1 I$ r1 y  A               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
1 r* M+ x3 ], J' x  Z( F* k2 m
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
4 n3 _  ]6 ?! \for i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
- S. k& L4 ^' E. a1 b1 N- E    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

- \* W0 H, f0 {2 _$ m* |8 \! u# 构造标签矩阵
. c' I3 w1 I5 Udef voc_label_indices(colormap, colormap2label):
3 |0 c5 @$ m) \) {    colormap = np.array(colormap.convert("RGB")).astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
( N0 i$ M0 y1 }+ g) a
可以打印一下结果
, {4 M7 S" z. o$ m$ M# d8 l. w
y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据

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

" O: ]" J7 {0 Z
def voc_rand_crop(feature, label, height, width):
    """
    随机裁剪feature(PIL image) 和 label(PIL image).
3 [7 |$ H! t( y9 l    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
7 ~3 Y( A) B! B, J8 D1 R7 [    Get parameters for ``crop`` for a random crop.
    Args:
        img (PIL Image): Image to be cropped.
# O+ o  s: I: c2 L; }7 k0 ~        output_size (tuple): Expected output size of the crop.
    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
) E1 e# S6 d, v4 f# L    """
* ^* B; @( r. @& K! l" b# a    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
$ p5 d) S- q6 w- D* c: q, v8 R    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label

! F, \* x7 `$ T6 q3 q, W# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
/ {, k" Y. w/ x5 E' F& Kshow_images(imgs[::2] + imgs[1::2], 2, n);
1 p7 h' `$ p: ]: c
% k  V' W  a" a


) p; H2 p$ ]' @
! t* ~7 O$ t- r: u- V
3 自定义数据集类

" [0 V0 P: q5 ^+ `6 B) B/ {3.1 数据集类
- I, k' i4 m. W3 q0 P. x
torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

9 S; h) W5 K  [, P__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
, ^* {; E4 l, f6 }由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
, a; M2 P* C+ [$ L$ D) e* D6 l
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
" U* g4 b! _( F+ s- R* i' O% c        crop_size: (h, w)
" v3 X3 G5 e) X  [' X0 y2 R# b" }        """
0 }) [- W1 U" A# C. ]        # 对输入图像的RGB三个通道的值分别做标准化
' G( e/ o, s$ [  ^1 v% r$ l        self.rgb_mean = np.array([0.485, 0.456, 0.406])
3 x6 p1 E8 o3 ]  a( c, C+ }. U        self.rgb_std = np.array([0.229, 0.224, 0.225])
# u) M9 G! I/ P+ i6 M        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
* Q5 h# l- D: _; p2 m            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)
; |! Q) c7 T8 t# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
! S6 `9 l5 a8 q& Q9 X2 g. e, F        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
$ I! z' q$ o; J* r
3 V3 a, N& c; @  \" u" ]; w    def filter(self, imgs):
        return [img for img in imgs if (
) g, q( ]. y1 }& C            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
' M) ^1 R/ j. B
    def __getitem__(self, idx):
6 s) P, [; C; F2 @( A# m' `        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))

1 e9 @) R' x8 g, @4 O    def __len__(self):
        return len(self.features)
# i3 {2 S( c. z; c0 v3.2 读取数据集

通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
, r) t; @; a) x& k$ ^* }
, [0 s' v+ b) y; H2 S! obatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
) {5 [  ]* R6 ecrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

# 创建训练集和测试集的实例
+ p# k' y: K" m" F% i/ Uvoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
7 w7 G% }  J8 [1 _; Ivoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
3 r4 _4 y' a( o0 g* q4 [+ E- Q
3 B; K  @3 }( d6 x" [# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
) K  Q4 g! c) f/ }4 Y, [* `num_workers = 0 if sys.platform.startswith('win32') else 4
0 h9 g4 k4 N6 G( Y. R) k' `train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
6 H* c' `& K' V- u/ q% Y( p                              drop_last=True, num_workers=num_workers)
. F/ d% \/ E8 V. ^) qtest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
0 I; Q" @4 S+ w8 w. j                             num_workers=num_workers)
: D2 Z* n6 F# k5 a9 O
# 方便封装，把训练集和验证集保存在dict里
  R% n; S# H, a+ h' U& }/ sdataloaders = {'train':train_iter, 'val':test_iter}
+ N* p4 v; D/ G" q& i" ]dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

1 Y2 s3 Q; W9 l9 w' j  Z+ I3 q1 l3 N4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
$ M( C% ]2 b/ N) a8 R
  F, @% |7 {7 V$ z" g- xnum_classes = 21 # 21分类，1个背景，20个物体
- t. y# S6 ^" N4 [model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
for param in model_ft.parameters():
    param.requires_grad = False
4.2 修改成FCN

全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。

) B" U) l' Q8 K( R$ g6 N  j对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
$ L5 x  |* ?/ e% ]
可以先打印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 倍即可。
( ~/ ?4 b, t- C: e! f- x% W
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
              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) # 转置卷积层使图像变为输入图像的大小

# 对model_ft做一个测试
# e3 ]- [1 K6 V* E0 q5 u9 Sx = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
( [9 ~+ {/ X4 ]+ [& @4 D
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
+ W& g9 }7 x$ |#     break
$ A; o2 e' S9 T: c. X  O6 O
: G# T0 @' I+ [8 E* K, h' s  D( Z: S. s
4.3 初始化转置卷积层

/ a0 H# b6 u, a. s/ Q在图像处理中，我们有时需要将图像放⼤，即上采样（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函数构造的卷积核的转置卷积层来实现。

( q+ w- C3 z5 U8 N! Q0 L
/ m6 }+ o8 m8 B6 [# Z+ ?# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
+ R0 A7 ?; u9 u- Q! @    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
        center = factor-1
    else:
        center = factor-0.5
" D2 a/ g/ V* l! J9 H& O    og = np.ogrid[:kernel_size, :kernel_size]
    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
4 {  y1 H. S) T    weight = torch.Tensor(weight)
, d. ]5 A  k0 L  u- Z0 f! b    weight.requires_grad = True
    return weight
% B% M; H" }/ K- g$ N  M
) h2 f  N6 b3 c( O
6 G* z  }6 G( \; I. P在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
) Z3 b$ q  g) X# q+ @: @# q! o  L( h7 @8 a
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)



8 J2 v0 @' m3 C5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
3 ]. s2 |7 w* @+ Z4 O/ t3 I    since = time.time()
! C- Y+ O3 C0 L/ [" r9 x    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
. k, V$ w5 |  i4 i& e0 Z: B' q    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
8 j4 `6 u7 M9 z) e3 O        print('Epoch {}/{}'.format(epoch, num_epochs-1))
/ x; H3 I. n" {) k/ ~$ N        print('-'*10)
9 @( i7 u2 v9 }6 f( u% \        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
, Q$ V4 U* t5 m9 G: I9 Q  y                model.train()
            else:
                model.eval()
            runing_loss = 0.0
            runing_corrects = 0.0
4 L( L9 c$ b  p; I: Y1 H            # 迭代一个epoch
( ^" C5 M3 m$ P. j: g! i2 Y            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
: ]4 c# J; j7 w9 d0 q# L5 j                optimizer.zero_grad() # 零参数梯度
( v  y6 l4 n' b( u8 y                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
% U$ Q+ T; Y7 M. h0 B                    logits = model(inputs)  # [5, 21, 320, 480]
# Q& d8 }) H- L& L: N1 u# G                    loss = criteon(logits, labels.long())
3 c+ |% t- n  v& t7 S7 J* ^) c! T                    # 后向，只在训练阶段进行优化
                    if phase=='train':
7 K" e- p- A: h) I% N                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
& a, w- c9 R" D9 k) K                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
, P% I2 c8 a; A4 x
            epoch_loss = runing_loss / dataset_sizes[phase]
5 n- d2 ~  \; M; J5 ~            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
+ N  w  D  t$ r5 u, X/ ?  k, c            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
/ i7 K! i5 b$ N. M2 h                        # 深度复制model参数
3 s9 B' V- R1 N0 u            if phase=='val' and epoch_acc>best_acc:
5 r; o0 v) f* n+ J4 n1 ]                best_acc = epoch_acc
/ ~5 F6 J2 V5 ]7 [                best_model_wts = copy.deepcopy(model.state_dict())
! |; }- I9 f2 E/ u( s        print()
    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
3 p3 ]' M: w& B2 m    # 加载最佳模型权重
: m9 D" S: B. w1 c    model.load_state_dict(best_model_wts)
* \$ D5 d3 N3 p% E* m    return model
8 ?# A5 ]7 l9 u& h. X- a5 |8 _2 ?7 f
6 t! f1 y( j$ f+ n% J, `下面定义train_model要用到的参数，开始训练

epochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
1 }9 j# H8 Z3 C% u8 M7 J' G. Joptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
1 b+ F2 j1 P, h. n7 x* y# 每3个epochs衰减LR通过设置gamma=0.1
3 |2 S; [( r3 P5 W+ Z" b4 ?: `5 Rexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
& t: P. R, i8 j: w6 p9 T: z
; p5 R/ z) U4 ~' u* `7 l& k' h# 开始训练
4 b+ k* P9 B+ b: gmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

5 E( V: T$ w6 t% o6 测试模型

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

def label2image(pred):
    # pred: [320,480]
. g' v9 I. [3 \& @, t/ Y+ D    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
$ o3 X7 q* L! p) i, n: N    x = pred.long()
1 H. {  E- r$ Q/ n- G    return (colormap[x,:]).data.cpu().numpy()

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
3 T. D' s5 B, j! {1 E& [4 W0 Qstd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
def visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
9 p1 q* F, |' Y- T3 Q) _/ y    model.eval()
    images_so_far = 0
" s( N0 J$ J( `2 [9 h, ]7 g    n, imgs = num_images, []
5 k  a9 d5 o2 a5 w' e6 K9 \    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
. \" c$ l1 c9 c2 f+ t$ m            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
+ P: h3 W3 g2 v4 J            outputs = model(inputs)
: q6 Z2 c0 }8 h" Z4 Q            pred = torch.argmax(outputs, dim=1) # [b,320,480]
! k$ Q; K8 G" J; Y, ^% J4 n            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
8 `  a% Q0 v- m6 f
( [: [6 ^% J' E& [            for j in range(num_images):
% X# v& x( M9 H: x                images_so_far += 1
                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
' Q7 j1 _5 |; H, w* _- R$ ~                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
                    model.train(mode=was_training)
  j, G. _$ o+ k% r                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
) I9 _9 n! ^3 ?6 a. g                    return model.train(mode=was_training)

" [) d( e* Y$ R/ w' v7 O& [# 开始验证
2 ^! R8 h+ K* R7 B( Q, P4 l7 @2 fvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
6 A" H. D! W, f, Y2 Ndef predict(img, model):
    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
% _$ E# }& x5 m% [            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
: b' J( b4 f, X& L    x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
" `' l) Y( ~0 a    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
0 y  A  ^3 ~$ B4 \  r$ t
5 {( o( Z- S3 b( g2 e# cdef evaluate(model:nn.Module):
1 w: [7 j; V! T    model.eval()
$ u+ R3 q& }# g    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
    n, imgs = 4, []
' v2 f- m6 G  D& b9 Y2 Q: u    for i in range(n):
9 p! K% Z) W* s$ `" i5 Z        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
/ c+ v& O. p) S( R) y! F6 M; M        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
; Y4 f0 _1 @$ J$ Z, @$ }
# 开始测试
. T: q2 e/ |/ Bevaluate(model_ft)

: ]3 B3 k1 p6 b, D0 a5 K/ |7 结语

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

1 b% ?9 J) s7 wEpoch 0/2
- g& o) R; H/ c, t* T----------
# ~: Z  Z7 `' g* Ftrain Loss: 1.7844 Acc: 0.5835
/ _, z- }+ W" X6 R  aval Loss: 1.1669 Acc: 0.6456

% R& g' N9 ]% l# rEpoch 1/2
----------
& N, m9 `* v. Vtrain Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929
3 |+ W  C( n4 j* d/ T
Epoch 2/2
----------
+ F% I9 o( Q, L. L; X0 itrain Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948
4 C5 D& }* |/ h' g# ?5 L; A
- J" Z8 }1 l, g. X5 M6 G) M0 PTraining complete in 6m 37s
% t' C3 }3 \3 r7 v; S


- |0 t( O( D/ {; A% s' D  W! `
; F7 T: c! r* d8 M9 b8 T3 S3 e0 O当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

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

语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
- B7 f- ]' ~5 b; e
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
& j, T5 h( P9 sU-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
% ?# ^7 A( z1 Z9 b+ Y还有（DUC，HDC）、PSPNet等。
2 F* E" y; w! G9 I+ X
常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

6 O6 C. R5 X- b8 z/ T0 b; n2 H/ n对于损失函数，除了交叉熵误差，也可以用这些：
+ M2 Q: c! J, c' H! H
% L- e! q6 h, A% ^Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
( |8 Y- S9 C( `& [% mCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。


0 X- o1 x& X% U, {2 T* ^: \
: @# z: H- I  y) o" g
- U' o9 t1 e. k! y; i
% b% ?2 ]5 r+ Y————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
) G) q" S$ i8 R) r/ i. p: A! ~, B% t

3 O6 G9 b1 n+ v