Pytorch实战语义分割（VOC2012）

杨利霞

Pytorch实战语义分割（VOC2012）
) u  d+ R/ Y! r& {本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
" ]- C+ a- [! ?
2 c5 p: W3 k+ {. L+ B* ]; d( e& z
语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
- ~- F2 s: X4 w( B5 q) l. L2 Pascal VOC2012语义分割数据集
* k2 y4 {& N0 V5 h- q; Y3 l2.1 导入模块
0 H( j2 D& G- x+ o8 N2 a4 |2.2 下载数据集
2.3 可视化数据
2.4 预处理数据
6 D% E) D; o! u5 J. j4 }3 自定义数据集类
3.1 数据集类
3.2 读取数据集
4 构造模型
' _+ y: [$ t) ]5 K) X  K$ B4.1 预训练模型
9 Q# G1 v: ~! `  y3 |) x4.2 修改成FCN
4.3 初始化转置卷积层
( e4 K- h8 \) ~5 训练模型
. i1 h$ h& E, B7 r- x6 测试模型
9 k& I6 }1 \0 S$ ^$ l6.1 通用型
6.2 不通用
; F" D) v# W6 D1 X* V8 {7 结语
+ F4 z3 q) }5 c+ k% a; j1 图像分割和实例分割

. c% z0 f. Z3 e( A计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
; ~- i/ d9 c9 ?7 I实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

1 E0 j) v' f" W, h6 V) Z0 P0 s$ X2 Pascal VOC2012语义分割数据集
' w  q8 A( N$ s  f7 I
2.1 导入模块
import time
/ j/ E  w. j) z+ E7 kimport copy
import torch
0 L7 Z1 D9 t  S* C# O) S' b6 a9 sfrom torch import optim, nn
import torch.nn.functional as F
( z$ k, _6 k$ T$ I: m8 q8 kimport torchvision
from torchvision import transforms
* H7 I5 y3 E+ j3 Gfrom torchvision.models import resnet18
import numpy as np
from matplotlib import pyplot as plt
* F- J/ W' U0 H( r& O  n8 Rfrom PIL import Image
0 @! O8 i; R% \import sys
# _) L; O+ }+ [0 |# rsys.path.append("..")
from IPython import display
. g2 U9 Y3 W! i5 J, I+ Ofrom tqdm import tqdm
9 J; E/ b4 l5 Z8 p5 A8 Vimport warnings
warnings.filterwarnings("ignore")

) w5 W% H0 I, I' |# b7 ]% ]2.2 下载数据集

语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
  B2 k; c+ l  {. P& z

ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
. W* ]( ]- d  Y2 r& XJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
- U& V/ q8 l# W+ p# @2.3 可视化数据
0 X$ x- S4 Y/ O/ g$ O. Q
定义read_voc_images函数将输入图像和标签读进内存。

def 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')
    with open(txt_fname, 'r') as f:
        images = f.read().split() # 拆分成一个个名字组成list
: C& d8 E! ]& _9 [4 F/ e5 k    if max_num is not None:
4 d+ _, O/ x% y2 A9 Y9 d        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
; e) Q$ X& j7 X% c( A; [' z- O    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
3 ]1 B8 {: D% f1 l5 v, u9 G* N/ ~# I        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
/ o+ |3 R) C5 J! C) Y8 ^/ ~5 e3 q
定义可视化数据集的函数show_images

# 这个函数可以不需要
" E- c' o$ C% G! W/ f( l# ?def set_figsize(figsize=(3.5, 2.5)):
8 |# _7 [# p9 P    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
- `: U0 L2 y4 L3 A( b    # 设置图的尺寸
; t+ Y% ^0 d7 X    plt.rcParams['figure.figsize'] = figsize
0 L/ h" n! \- e+ s5 n4 \
def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
! }" J, D9 }& l    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])
" A$ o1 h/ J/ E" f            axes[j].axes.get_xaxis().set_visible(False)
8 S0 |7 l. p2 G% Q5 _8 F            axes[j].axes.get_yaxis().set_visible(False)
: E$ ^( P1 p# J% ^7 y5 B    plt.show()
; `& A' D  Z( y% o    return axes

定义可视化数据集的函数show_images。
$ X+ ]. V& z, {( ^* V( B/ k
# 这个函数可以不需要
9 R9 Z& ?1 y8 h, J9 T% r8 r& ?def set_figsize(figsize=(3.5, 2.5)):
0 |) ]$ T5 _& b2 w  k7 j3 u    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
7 h; {% ~" P2 Z    figsize = (num_cols * scale, num_rows * scale)
7 I; r. g  H; G/ h    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
        for j in range(num_cols):
( ]  Z  {! u7 A+ I7 ~            axes[j].imshow(imgs[i * num_cols + j])
5 I. M& u- B6 L9 o$ Y            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
5 M3 }% o4 v( y画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

6 u, Q1 T7 t! s: A5 z# 根据自己存放数据集的路径修改voc_dir
4 K& f6 F4 h+ }0 f( O" e0 z) ^voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
, t, Q/ Z' z# N' g9 Q) t  Aimgs = train_features[0:n] + train_labels[0:n] # PIL image
: e2 F. O3 y+ X- D, vshow_images(imgs, 2, n)
) R5 ~1 S2 W  O/ [/ Y
- h# \. U- R  w. z1 j+ E
3 d0 l1 i! @8 T& f. C  t
. t7 l% \  A* A& a2 _* _列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个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]]
8 j, W! m2 U4 F& E! p$ }) n6 b* p1 g# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
. R1 G5 R8 y, y- r  y6 `  \2 A; O有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
' v) p, n6 \& c; N8 a* m
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
$ d  t4 n! C3 ^# O! Hcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
7 b% i+ U4 j5 c" k3 ufor i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
2 m$ b4 ^6 u0 Z+ A& I/ P9 |4 n    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
, G$ @* F; u1 v1 Z' d4 c# l* g
# 构造标签矩阵
0 d/ Q( z8 p& T0 H( Ddef voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
/ T& Y; C9 J& l1 m0 S  i7 j! V    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
6 o9 }+ x1 ?1 `$ F6 z9 |    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
6 y1 ~, Y" L- V, \! D
7 d- x7 G4 P8 N可以打印一下结果
, z1 }- o$ w8 C2 [: D' T4 Z$ e
& }# V! A8 P% E  Z) X: e! G. cy = voc_label_indices(train_labels[0], colormap2label)
2 Z  ~: N! a) ~9 F/ P5 mprint(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
. O: n( Q5 |( f! {
2.4 预处理数据

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

8 r) o  Z% w. a  O. }
def voc_rand_crop(feature, label, height, width):
    """
. X9 |0 ?0 r5 Z! B5 p4 w    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
" T7 V7 ?9 }1 _" R# M) _    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
# A! M& C; ^. b  t# g    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
" f2 E: v0 ~3 D) v    label = torchvision.transforms.functional.crop(label, i, j, h, w)
' F+ t/ _$ t- S    return feature, label

/ C1 h7 j  ^2 t- b: ?) p# F# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
7 W$ x7 u8 W( [  J    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
) h8 j: e/ m) `, e- A$ _show_images(imgs[::2] + imgs[1::2], 2, n);
$ f2 S9 z* e/ B0 \5 Y, `
2 i( E8 ]  W% b" G  a/ A
2 l, R" J, m9 G* k9 c6 Q/ u



3 自定义数据集类

+ r" d, }/ ^1 J. U2 q5 Y& @3.1 数据集类

! @6 s. x2 i; w# p" Q3 d. d9 etorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
6 |  a  Z4 Q8 O! o
__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
" s5 O& V3 P* v/ e4 O, p8 y/ E8 o. ?
) x) k% U$ Q) N- r: @3 n, M$ s* l2 Yclass VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
' y' K% ?! \6 B/ |, q        """
# N% L  l4 j! x& S! R        crop_size: (h, w)
* u8 `7 Q, [+ Y        """
        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
. _0 Z5 ?1 Y2 O# Z" K2 \, K8 }        self.rgb_std = np.array([0.229, 0.224, 0.225])
        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
5 [. o3 n- E/ p+ @6 O3 y% R0 |7 u        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
+ {! H; k2 d: G0 |2 h& I; r# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
/ C3 n& _+ M. [8 d" w; F1 j9 z        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

! _8 U0 C. F: C. {) p; ~4 _5 O. O    def filter(self, imgs):
$ s6 W0 ?/ X+ Z$ r        return [img for img in imgs if (
* S$ R: {3 h0 ~. a5 L) m            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

& G& M8 m8 \2 d7 U/ g$ Q    def __getitem__(self, idx):
2 b1 q, O, u* S* b$ A: _" l        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
4 @5 M- r3 A. A5 m8 W; c( E                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

, C5 P9 Q( r% ]7 m, T    def __len__(self):
7 _/ Z1 y6 m2 C6 H6 |$ x        return len(self.features)
3.2 读取数据集

通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
1 w6 ^/ [1 D) e3 h
1 ^+ _% t+ m1 Z+ O5 ?, b1 V2 Wbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
/ F" Z5 D+ e$ a  @$ P0 avoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
! R: W6 a1 N4 unum_workers = 0 if sys.platform.startswith('win32') else 4
) I6 ~1 b: W6 F3 O) y! B5 E% Utrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
- m0 E* {* y7 R1 v5 V                             num_workers=num_workers)

$ J! i# o( A- E, W" a# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
9 O8 Y+ B/ C) Tdataset_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')

2 \2 C% d1 j5 j5 S/ s* i8 p$ N. W% rnum_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
5 J$ y" S$ X5 m  ]# f5 v% jfor param in model_ft.parameters():
    param.requires_grad = False
4.2 修改成FCN

全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。
+ z& F9 |" H% }6 }, q
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
$ P  d* J2 h9 e/ v4 t
可以先打印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 倍即可。
( L' ~) E+ h2 ?. {* `% j
  r  Z0 X' N( d' ~3 ^9 J8 omodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
0 h! ?' U! G4 h              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
" g4 A8 h! u$ x3 j/ z: g              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
2 I( `: H7 D2 n+ H. A+ z' Q$ u
# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
: I8 n8 x6 e+ s# `0 B/ d% ]print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
' r& \. V' t: P# for X, Y in train_iter:
: Y1 X+ p# H; m7 h. K$ e. V#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
#     break


5 [: P) h5 g6 T6 w, ]' T! p4.3 初始化转置卷积层
9 `2 R3 E* @! n3 B7 U3 f
5 ^+ H( I. [0 o) Y在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
6 n2 C& o* p% O( p# t; T# y在坐标 (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函数构造的卷积核的转置卷积层来实现。
  N9 J1 c7 y* R, P4 `9 G

# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
- ^6 W" ]4 `( N! K) t  h- A3 l    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
0 \8 A* E, B6 r; q3 W        center = factor-1
    else:
        center = factor-0.5
    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')
$ v1 v. W4 x# x# I, {4 R    weight[range(in_channels), range(out_channels), :, :] = filt
) S; W- ]! y5 Y0 M    weight = torch.Tensor(weight)
    weight.requires_grad = True
+ g, j. l7 M7 N    return weight
% i, e% Z% D* W/ }7 M) N4 `4 K/ ?" _! L
2 u* t- G/ e, K; X7 ]
# |. x! G+ G0 k+ y在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

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)


, u3 Z! W  `4 T: E6 C3 k) }, g
4 Y5 U7 o8 ]. Q1 X, i3 p5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
0 Z, e" j8 H) t: z    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
- l; F5 f2 X4 `4 e! t8 Q    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
& A' ]) `( s% G& ~% G! b; E        print('Epoch {}/{}'.format(epoch, num_epochs-1))
. x0 A% e$ ~* R& e% q% t$ k        print('-'*10)
        for phase in ['train', 'val']:
5 Z; D7 ~) D8 Q3 {: m4 n# J            if phase == 'train':
                scheduler.step()
4 z2 H2 M$ E2 N% m                model.train()
            else:
  g& U& `" m0 N5 _3 \: t                model.eval()
            runing_loss = 0.0
            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
1 A: t7 l7 b! `6 h9 P                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
# y0 y- J) L' [; `8 |/ t                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
& Z5 P# |' M) e4 @% u& C                    # 后向，只在训练阶段进行优化
                    if phase=='train':
                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
4 W+ j5 y+ b0 y! S8 d. p                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
7 h$ g, P9 J$ o- C$ H
, X0 T8 l4 n+ ]8 q; k$ q            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
' d6 e# D! U" m$ P$ \        print()
: `# v$ f( m9 {) G    time_elapsed = time.time() - since;
0 h8 g- [0 z- t+ J- i2 d    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
1 H  Z$ [2 g5 m$ o* n    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model
+ A# I$ _9 q6 W9 y: P; E/ Y1 ]
下面定义train_model要用到的参数，开始训练

9 x, b, n$ ~8 M  m/ V+ r5 }7 Tepochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
9 X9 @; m0 c4 f  b+ Joptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
4 K( S1 e6 a7 e- G3 `8 \# 每3个epochs衰减LR通过设置gamma=0.1
8 h8 a9 [' O+ J  Bexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

; C4 q  P( `; P! @# 开始训练
% L  }1 W* L1 d' C9 N. d. V" ?model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
6 N8 ]& f3 k, j" I, ]
6 h" _  J+ x9 {6 y/ B' j' q6 测试模型

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

def label2image(pred):
    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
: r0 i4 I6 Y. f# |+ O; X# @9 f' b    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()
% n3 \. l4 U  [, Z1 b
3 X, ]- S- \  N0 u$ @

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
& m/ W6 {8 m" R  y( I" Estd=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
6 u* u+ R8 i! B: c+ i; U* @    model.eval()
    images_so_far = 0
5 j$ A; M2 ], i! I% X( ]    n, imgs = num_images, []
4 F; n: t3 Y& @( l, Q& i- U    with torch.no_grad():
- ]& H. O! a0 l, o4 h( _" Y        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
& f9 c- F& X* u+ N            outputs = model(inputs)
            pred = torch.argmax(outputs, dim=1) # [b,320,480]
8 C1 u0 d7 W' W3 n9 ]9 n0 P            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
2 S& X2 S' Q) G+ a7 J
            for j in range(num_images):
' e+ ^. K" _& m" c: j* S3 Y+ T                images_so_far += 1
6 A( b( x7 c  A7 O  F                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
7 K$ u5 R+ E  v9 p9 O                if images_so_far == num_images:
+ P3 y0 W  N  ~/ f0 m! X9 A                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
. T+ x8 l* ^6 |9 ^( A                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
6 i$ B% c( o0 ?9 v6 u                    return model.train(mode=was_training)

  s: P8 F8 X6 N; ~# j/ w' L3 F# 开始验证
" P4 h4 b4 g4 Gvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
( w( m* b# N' t% |) c6 ~def predict(img, model):
    tsf = transforms.Compose([
: m6 o: [7 v% T( I, L8 a            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)

- I/ Y2 o8 T- f+ Rdef evaluate(model:nn.Module):
3 j7 i4 k" n- q  j5 M: b5 P    model.eval()
$ b! x9 u6 H% w0 L& j( Q    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
& k9 D4 a8 n, V0 ^/ `3 o    n, imgs = 4, []
    for i in range(n):
: h/ a, R  [0 z/ w- P5 x/ Y- N$ ]        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
( Z& T$ [, \: n        pred = label2image(predict(xi, model))
0 ^2 |7 V/ R) O* A( w/ h        imgs += [xi, pred, yi]
0 k) K# H! }/ Z8 `! T* ^    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

) O+ V# k5 v8 d8 z; ?# 开始测试
+ k, a2 k7 l5 r* y/ _evaluate(model_ft)

! ]8 @9 r# e* u  a) g2 a- `; a7 结语

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

7 b9 ~4 y3 Z' n# C8 ?Epoch 0/2
' H, Y1 H' P0 \, l7 O4 @----------
3 c3 T% o7 g! {) q+ x5 Q7 L6 Qtrain Loss: 1.7844 Acc: 0.5835
( }9 J) u! ?7 N" K; E) Hval Loss: 1.1669 Acc: 0.6456
, [$ M, X! T1 z: s5 a7 {9 y- x5 t
  _2 s+ ^2 y3 @$ _Epoch 1/2
----------
0 V& v! w; s( B3 c( strain Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

Epoch 2/2
+ R/ `3 D6 C4 C8 n. v" \----------
train Loss: 0.9578 Acc: 0.6706
, l0 X# s( m) n9 B5 V/ `. uval Loss: 0.8088 Acc: 0.6948
3 @4 ]  W6 c( L& R) q4 _% d
Training complete in 6m 37s


. Y5 r( }) C, @* x
) A- ?$ }+ n: V2 d
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

3 z7 X2 C2 Q/ B( O0 i对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
! M# `2 }. d5 ^8 z( e7 }4 M! K0 L, Y

语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
0 I' u* \9 a: i6 u: y4 J$ `
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
! O! Y+ N: c. EGCN 通过全局卷积网络改进语义分割[论文]
9 C2 I+ k& r8 j! f2 eUperNet 统一感知解析
3 ~( o; V# ]. [: `9 i) YENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
$ i: s: [# `# u; a7 j3 eSegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
$ Y. `  A+ ]/ e8 g' x! A# Z
0 w6 P) e! ?8 f) R+ j2 _对于损失函数，除了交叉熵误差，也可以用这些：

, t. S. J! F7 }+ t8 k' nDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
+ B8 W6 ~# {# D
# _# g  L0 g* A& q


4 @! w# B5 ^$ E9 k+ O3 j$ d
4 C( ^% J% n- D4 o) H————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507


: K+ c7 v4 F: h5 n
! y$ I9 \: r/ l& U1 m