Pytorch实战语义分割（VOC2012）

杨利霞

& r1 m' \( ~& t, c
# Q- ~& g, e- W) JPytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
4 y6 m0 Z# P( e, j  l语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。

; i( B) |' f% h5 b4 J
语义分割中图像有关狗、猫和背景的标签
+ T8 F) i7 M8 b8 y文章目录
6 L! I* m) {* M' _
1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
2.1 导入模块
! w5 z: ~; y9 E; ?$ e; l2.2 下载数据集
% `8 N3 y3 K4 r% [$ Z2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
3.1 数据集类
3.2 读取数据集
4 构造模型
' S0 r7 r5 E3 T9 }4.1 预训练模型
# R) _$ N8 [' o1 j4 M+ I5 I( T& I: R4.2 修改成FCN
4.3 初始化转置卷积层
9 J/ G  e7 U1 e; q6 y- c/ o5 训练模型
6 测试模型
6.1 通用型
6.2 不通用
7 结语
1 图像分割和实例分割
5 K* ]0 ?: N5 J
9 a2 X/ e6 B' Y/ z5 |$ a计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

+ p2 R( }* H' X8 Z' o图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
" [! D' f. Z0 q* ?. @
( a2 o9 w9 ~; K- y2 Pascal VOC2012语义分割数据集
1 B, h/ h' {4 d2 ]6 |
6 g( Q# ?/ B4 L; t( b5 y2.1 导入模块
import time
( z3 L+ d  a) h% n- Limport copy
# v: U+ Z) M; {7 _* t/ J: vimport torch
from torch import optim, nn
import torch.nn.functional as F
import torchvision
. K( H- Q0 F0 W) F7 M4 T% sfrom torchvision import transforms
from torchvision.models import resnet18
import numpy as np
/ I- w2 d! q  H! X; Kfrom matplotlib import pyplot as plt
* v% V  p, k, l! efrom PIL import Image
import sys
sys.path.append("..")
; q! X' _' y  W9 wfrom IPython import display
$ N  F" _5 a' h2 I0 M- t0 Hfrom tqdm import tqdm
import warnings
, T8 ^( ~4 _# X! p8 p: g  Y# D# Xwarnings.filterwarnings("ignore")
6 C* I/ }, Z4 g5 {
+ u( J/ f' w; |4 U2.2 下载数据集
& i1 e% \/ O4 W4 x
语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：


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

定义read_voc_images函数将输入图像和标签读进内存。
  @; A! u, f1 ]6 u
7 S- L: k! ~" Y% k7 ndef read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
1 w2 S; p# Y/ X: ]/ c1 b! G# l  o    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
4 W" `/ ]3 T6 K. n    with open(txt_fname, 'r') as f:
        images = f.read().split() # 拆分成一个个名字组成list
    if max_num is not None:
$ Q6 e2 V4 a. H9 x        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
3 w4 h+ S) C! m# ^2 E        # 读入数据并且转为RGB的 PIL image
' }' J8 s$ v' Y. H( M8 B$ F' @        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
. z0 U$ H  w; K8 Y9 C- M        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255

定义可视化数据集的函数show_images

7 L( I' @7 S0 u' H+ X; ~# 这个函数可以不需要
2 g+ x% V* G/ B5 Kdef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
: C5 o9 U" H8 ~    display.set_matplotlib_formats('svg')
/ p) x8 Q; O* ?: n  s    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
- e1 w( b# Z% _* s, @) A
* u# i. A$ D6 }5 b$ t; h; Cdef show_images(imgs, num_rows, num_cols, scale=2):
* A( K1 c8 m; ^6 |/ Y- Z8 V" R    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
  F: e) m5 S7 J6 a! K+ S3 X    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
        for j in range(num_cols):
, b/ @" Q6 e& ^            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
' ~  ?3 k: [, w' k- F/ Q5 u            axes[j].axes.get_yaxis().set_visible(False)
/ a0 ~# F! Q2 D" {9 R- H8 ], w    plt.show()
( R' c, d; {# c% A    return axes

5 f+ \5 ~4 Q' R7 w4 W- v定义可视化数据集的函数show_images。
. z5 I# \, @6 ~* u8 X7 j3 ~/ F
( d7 ]0 t& v+ a0 q/ H$ @9 b# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
% U2 U+ t# q( x8 p    """在jupyter使用svg显示"""
  o" g" s$ E+ q2 L5 S7 d    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
2 f$ D0 T# o0 Y# y% n6 w    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
+ |' U) u* e2 k# ?. D    for i in range(num_rows):
        for j in range(num_cols):
7 }9 s% P/ p' \4 o9 z            axes[j].imshow(imgs[i * num_cols + j])
7 F& D5 h7 p  ~' m; j            axes[j].axes.get_xaxis().set_visible(False)
" |1 k. @' X& u$ D3 q1 B: v( u            axes[j].axes.get_yaxis().set_visible(False)
& ?2 l/ J9 |% {8 T/ P3 |    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
train_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)
0 I8 i! O2 }2 N. @( `2 f# b


  a2 C/ N; V; ^$ k" Y, V3 O列出标签中每个RGB颜色的值及其标注的类别。
: r% \- k( l# T; {# 标签中每个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],
# P9 C+ ]/ i0 k                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
& m6 T+ I& ?9 Y4 U3 d                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
5 L* P! a, f0 D2 E) I% S; L  c& |                [0, 64, 128]]
' d7 {# u$ S# c, f0 v# o1 T# 标签其标注的类别
! c4 F2 G+ S4 ?) v0 f; lVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
; T1 e! O) m+ ]  ]有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
9 r& U4 [$ U9 zcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
0 f# h) A$ g. c7 K: a+ Q8 m& Ffor i, colormap in enumerate(VOC_COLORMAP):
1 a3 e3 ^  k" P) ~    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
/ K- e6 `! A) e% s+ C1 Q    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

% U; p+ A7 Y4 T# 构造标签矩阵
" m' ?4 G5 B& P& O: }; ~0 ydef voc_label_indices(colormap, colormap2label):
! Q1 [) U. p# m. a; F) H; x. p$ a    colormap = np.array(colormap.convert("RGB")).astype('int32')
7 i  V" J  i; w' }7 K& ~* W) o5 d    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
$ P% a! Y4 J4 O! Z4 }: R0 n; i  o
/ V: v! I2 _$ O. T) e7 B9 y可以打印一下结果

y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
. M7 G- e, d# z" z) R, c9 A
2.4 预处理数据

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

3 Y8 ^1 ?% P* J+ l4 D2 C
  U+ s/ r1 l+ Y0 c2 t, fdef voc_rand_crop(feature, label, height, width):
+ H, B5 N- B; F0 a    """
    随机裁剪feature(PIL image) 和 label(PIL image).
/ ]8 d; X5 _# E( ]& j% W    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
  ~$ b/ s: Q6 G  N/ X' C        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
    Returns:
& o& ~1 R% r9 X, f        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
- |0 L; V: j& ~& G: `5 _$ I1 L0 V    """
    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
! b2 G+ I3 N; K0 D' ?    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label
" O$ \% @1 j7 L( d; t* A/ ~
# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
: J" W& w+ a* [( pfor _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);
5 G, e' I$ q" D" t) Z7 P


5 v- D# X! T3 k' A3 O


$ B! S: S- ~: d3 自定义数据集类
6 V# f2 o- t, J
$ |1 q: ]7 ?6 v! ]# j! i4 ^3.1 数据集类
8 a* V7 @5 [4 h0 ^6 }5 o
( b, I9 o( W: x6 d6 |0 ^; Htorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
% H% V6 L9 w. o2 N, X2 t
__len__ 实现 len(dataset) 返还数据集的尺寸。
6 q; ?8 u& |$ N__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
( ^) p6 L/ f3 q% C* e
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
! A' G; [: m' u1 {' }' M8 U        """
- f( T; D/ S" }  M& y* x7 o        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
3 c8 S: _  }7 B( h* o: T        self.rgb_std = np.array([0.229, 0.224, 0.225])
; ]  l4 u7 f- ], T3 i- r# s        self.tsf = torchvision.transforms.Compose([
# y+ ], i/ ]$ e6 B& p            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
        self.crop_size = crop_size # (h, w)
6 N  [" _8 p9 o$ n        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
; R: G( j$ U# ^) x& v$ g' g; R" `        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
- l" U9 H7 V' {8 Z. L        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
7 W3 H/ L6 X+ V8 C
% G$ K; W, _, k    def filter(self, imgs):
        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

2 I0 a3 c! @% |/ x6 V. q  n/ y: \+ b    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))

/ x- m( r; g" d2 x' u    def __len__(self):
. k. _9 e" }6 S: V+ f/ u" B        return len(self.features)
3.2 读取数据集
' Q( h* b8 ~  F* {0 j6 H
通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

% c5 `& x1 h1 v8 o/ Fbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

# 创建训练集和测试集的实例
" q. j$ s$ R8 t3 svoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
! V2 e- ^3 I" E# |4 ^6 \8 Uvoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
train_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,
! C* s. J, Q, u% `7 r; c0 i. k1 n                             num_workers=num_workers)

3 |! j! o4 P8 P/ f9 z! u% J# 方便封装，把训练集和验证集保存在dict里
* k7 W5 o2 v+ s7 E' edataloaders = {'train':train_iter, 'val':test_iter}
dataset_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')
8 O& h" z2 f3 d8 K2 d; k
! K  l0 p3 f; F& n- @. Fnum_classes = 21 # 21分类，1个背景，20个物体
" Q) ^; m" r- z7 `' J2 G% `0 ?model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
! \. b* b0 D" V" Y" w! N+ p
# 特征提取器
; Z* e6 U) S5 X3 d* @# u) Afor param in model_ft.parameters():
4 B1 K3 F" [: `% s7 ?    param.requires_grad = False
4.2 修改成FCN

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

4 z4 X' z2 L( Q' j- B  W对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

' U5 n; L" E) l, k可以先打印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 倍即可。
( p7 K% P- Q; Z* ~. i4 a/ v* ?' v
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做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
  b+ j+ v5 l1 S6 J2 k4 H5 `
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
% M. W. k8 W2 G2 A4 q* B% @#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
' l1 L! N" O+ C3 b( }2 N% h" \#     break
5 @% @) |" B' Y5 r6 o
. q7 }; M0 q% K  X* h
, ?/ @# u- M' V1 p4.3 初始化转置卷积层
) @  Y( Y) g) H9 G  ]- u
& m! R2 I% [( p. S  v& ]; \在图像处理中，我们有时需要将图像放⼤，即上采样（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函数构造的卷积核的转置卷积层来实现。
) ]9 S/ o) U6 e! R: y) t
8 x2 c$ ^- |0 D
8 R* p8 p; ]- {6 b& S' ?# 双线性插值的上采样，用来初始化转置卷积层的卷积核
5 A! l1 p' Z- G8 @6 J: pdef bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
. I% O3 w# z  ]2 g- O4 h5 O    if kernel_size%2 == 1:
6 H& {- Y' R, W- o. F3 F        center = factor-1
  p$ w; K& W8 v  Z! k    else:
& E, u6 R2 h4 k2 e8 W% }( j; e4 b( ]8 `        center = factor-0.5
% Z' H' ]. X7 b8 |6 q. R    og = np.ogrid[:kernel_size, :kernel_size]
# a# {7 i, T1 r0 j  A    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
    weight = torch.Tensor(weight)
! u5 C3 d; O/ A9 [: F! H. X& K! @3 _    weight.requires_grad = True
    return weight

; m" N: ~! T; Y
9 g0 ~: k7 I5 V在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于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)


& @, E- D* n" K9 v
5 训练模型

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

5 l* I/ O3 L* j7 i5 m/ N$ K1 Ndef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
/ p/ N& Z2 K, N# v5 N% o    best_acc = 0.0
7 C% I3 f  T" F0 q4 j( p    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
8 ~9 f) \* r$ E& d3 \$ k( o        print('-'*10)
        for phase in ['train', 'val']:
% p. p+ B. T/ j, c: ?& C$ z3 T% j: {            if phase == 'train':
                scheduler.step()
                model.train()
3 n  [* i9 p" Q) j7 B/ [8 ?6 Y            else:
' z& h" H) X. e3 R. t, {                model.eval()
3 Q! A! ~( _2 ~6 B# x; I' r            runing_loss = 0.0
+ b; c5 ]  Z8 z3 l6 V            runing_corrects = 0.0
            # 迭代一个epoch
) Z5 Y1 Z- T4 t- s) D4 Z" O            for inputs, labels in dataloaders[phase]:
  ]! m; z7 q% T3 U& t8 w  W" F  I                inputs, labels = inputs.to(device), labels.to(device)
. J: h: R4 l' M                optimizer.zero_grad() # 零参数梯度
                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
" i$ Y- ?! Z6 X. L( [                    # 后向，只在训练阶段进行优化
                    if phase=='train':
' w( v0 o7 ]' R8 B& ?" l" G- l( j                        loss.backward()
6 F' c4 U! k( I5 h8 [                        optimizer.step()
3 `- k9 J  i/ b" j/ R                                # 统计loss和correct
# h4 B1 X5 _" ~! }( M2 P                runing_loss += loss.item()*inputs.size(0)
4 L3 n: r& ]$ j- ^! N- B( [                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
' x1 B6 J# Z& O
            epoch_loss = runing_loss / dataset_sizes[phase]
  o& g8 j9 d  n% z# [. q            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
! l& d$ e  B# D% Q; ^$ Y' j& H                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
/ c0 b; N0 q9 C                best_acc = epoch_acc
  W$ O7 F) r9 M, o7 j. D, i& I                best_model_wts = copy.deepcopy(model.state_dict())
$ J( M, i) [: X, f* g! m" h5 Y! O        print()
    time_elapsed = time.time() - since;
2 |4 L) J4 v& S. Z    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
3 `$ P9 y* }4 W8 q& [    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
& p  }5 Z: m2 O) z5 w) e    return model

下面定义train_model要用到的参数，开始训练
, ]! \+ t3 D% S  P+ ~
  u  i4 w) B( \! M4 y! H- \+ A. Kepochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
2 H: Z9 [/ m- u; H+ Voptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
3 f' |% A$ Q$ texp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
" [! c3 W: q8 e% y7 D
7 e3 n9 m) X+ d# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
( W2 k& w7 m9 K* t2 h8 R) Z: {
, Y% \8 i! B, g. P8 R6 测试模型

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

def label2image(pred):
    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()

1 j( w( k. z% [) `5 M1 D. X

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
# g. N- I; S% I- y0 ]- edef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
0 o, s- ~& S# V8 y$ h3 W: x    model.eval()
    images_so_far = 0
    n, imgs = num_images, []
6 U% I; o* X8 w, h; K. G    with torch.no_grad():
. n' o* Y# s8 C& I        for i, (inputs, labels) in enumerate(dataloaders['val']):
4 y: c4 l# T3 ]3 f/ \) O3 e            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
# Y5 L0 @5 F  M2 ]% f* ~            outputs = model(inputs)
            pred = torch.argmax(outputs, dim=1) # [b,320,480]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
# C" G2 P: z/ Z4 A
! S1 [% |' N( E0 Y" g2 {2 J            for j in range(num_images):
                images_so_far += 1
: L8 K9 r6 V- e                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
7 M" ~7 [/ n0 z) O% ~, s. u                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
                    model.train(mode=was_training)
( }5 S' l+ h! O, K2 O/ w1 f                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)
# m( r5 @! n, c% }9 S# X* s
# 开始验证
visualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
3 ]; l! B2 Q: Q4 l- u0 M            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)
: }1 x: e. y/ }    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)

def evaluate(model:nn.Module):
    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
) [  `  R. v: T1 A3 N( q9 e0 X    n, imgs = 4, []
    for i in range(n):
0 _, `, S* W$ S        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
8 d, q4 D- w6 |5 n" T- O    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
6 k4 r9 }; r2 Q6 D
' z/ ~3 n* R5 y9 |# 开始测试
/ ?  P0 ~% A; m9 ]5 ~% Gevaluate(model_ft)
0 m5 }" N' v$ g# O( u8 L$ B- M' {
1 q# X) d# p+ N0 D1 Z+ O! h7 结语

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

Epoch 0/2
----------
# U* W+ K: h5 e* \) Utrain Loss: 1.7844 Acc: 0.5835
1 Z( U* t  Z. v8 K$ ?% q. @val Loss: 1.1669 Acc: 0.6456
% l9 l. T: b. o" `& Z: t, h7 w
Epoch 1/2
0 l& {, G' j2 A) D1 w" k& W----------
train Loss: 1.1288 Acc: 0.6535
+ L1 Z! _1 @7 a; }# D) ^val Loss: 0.9012 Acc: 0.6929

$ X- p/ h8 I% H: H: A: j; \+ u# m- ]Epoch 2/2
  A3 d6 }7 ?. Z: N----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948
. n7 q4 v+ f: `* o2 O
. ~' `5 b# V4 p4 J* P/ d* F4 {/ DTraining complete in 6m 37s

0 s3 \" ~! M9 @& H- E


, G  W; R1 R. z( E4 {( L/ j当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

, K% a& d+ `3 p) x对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
) ?0 W" m4 }# `1 c' M

$ t- q3 z" A! Q' q8 k1 `语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
! Z* c  U3 Z4 b3 T
- G' g; l) p- G6 m  ], @# bDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
# _% r/ G# w, L0 qUperNet 统一感知解析
& f: O" a: x) f8 MENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
( {, ?& g( Y: p" VSegNet 用于图像分段的深度卷积编码器-解码器架构。
; _4 G" i9 \; B$ d还有（DUC，HDC）、PSPNet等。

常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

! u0 k- j! D0 i$ D% j& n对于损失函数，除了交叉熵误差，也可以用这些：
( X3 p' Z/ j8 K/ B( O3 u, M# H% a" v4 R$ e) n
5 S+ K! o* D0 nDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
. C) ]: c7 A. X7 a4 x) I) q! GLovasz Softmax 查看论文：Lovasz - softmax损失。

8 X: z0 p9 r- B; L5 y2 x



7 X: T; ~6 u) v- {' X————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
, p( v; `7 `6 k8 g- s5 w原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
% @* Y6 I6 {6 G" i: l
- D6 @* C* z8 y& ]! ]