Pytorch实战语义分割（VOC2012）

杨利霞

* X' l! X* o9 q5 q! _: }8 [Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
2 {, W, W+ A+ Z' Q  D* @8 T语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
1 j; p6 J3 T6 l( g# W4 m
8 {' @* A5 x: V/ \: |( q
( T2 S! I4 n" ?1 i  I语义分割中图像有关狗、猫和背景的标签
文章目录
- k9 H3 ]( w! \2 l, F
( A% X5 `" h0 Y% d- x1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
0 [+ Y9 j$ k3 z2.1 导入模块
0 c" H6 q* z* |5 ]. ~  M2.2 下载数据集
2.3 可视化数据
8 N7 \! S; v0 f2.4 预处理数据
6 C# U0 w1 N5 A& A9 m" N3 自定义数据集类
. v0 \/ P/ v- z, w3.1 数据集类
3.2 读取数据集
4 构造模型
4.1 预训练模型
+ Z" F3 S6 y7 _$ k4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
6 测试模型
6.1 通用型
6.2 不通用
% ?# [: a% T! O+ a7 结语
. n9 ^1 k% h6 z& Y& k6 z: o1 图像分割和实例分割
8 ^' ]" s% K7 S. ?4 l
计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
/ r) e& }& @* e2 P! j5 T实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
" W$ \% ~% U: |% l- e
2 Pascal VOC2012语义分割数据集

- p7 X) e6 v" }% i2.1 导入模块
9 x8 C# a0 K7 B4 y# |/ Fimport time
import copy
import torch
from torch import optim, nn
import torch.nn.functional as F
import torchvision
2 n4 v  P, U: @+ p& Xfrom torchvision import transforms
- _9 k! }; A; j+ v0 N9 h2 Dfrom torchvision.models import resnet18
* k  o+ O: ]" e& E8 {import numpy as np
from matplotlib import pyplot as plt
1 c. T' L; T6 H9 ^from PIL import Image
5 s) @$ R5 F0 v9 T+ R' qimport sys
sys.path.append("..")
7 M) n5 `9 f8 b: |from IPython import display
from tqdm import tqdm
) a1 O: w. |4 y( {0 ]import warnings
: t% j1 @4 \# |3 Awarnings.filterwarnings("ignore")
- Q0 w' g8 }  T: Z( o
6 m9 E! a5 W- ~2 N! O! I2.2 下载数据集

语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
6 x5 r1 D3 F$ w& O0 P1 _) V3 I

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

定义read_voc_images函数将输入图像和标签读进内存。
5 `" h9 R' Y$ P) V
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')
6 b, e7 B% ]% W4 K7 x6 w    with open(txt_fname, 'r') as f:
, m+ ]3 u# U) `6 [: Z+ T        images = f.read().split() # 拆分成一个个名字组成list
" K9 _. `% b8 b$ y$ l# A! ^    if max_num is not None:
        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
- ]: [. [6 ?3 B% c    return features, labels # PIL image 0-255
. f' Y. p( Y0 l6 U% \9 _2 j
定义可视化数据集的函数show_images
2 z/ T1 }, Y0 W
# 这个函数可以不需要
# i* _6 `2 y" V: Wdef set_figsize(figsize=(3.5, 2.5)):
+ @; F8 m, @6 p1 c$ L+ W: i    """在jupyter使用svg显示"""
+ P- L4 x( Q4 f    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
- y8 r; W( P# z& p, w( `    plt.rcParams['figure.figsize'] = figsize

% r2 H1 s) C+ @# @; s$ Tdef show_images(imgs, num_rows, num_cols, scale=2):
2 g8 H. M- K/ Z5 f# ~0 l    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
# o! P& _' E, R, G( |* f    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
6 b! ~. ]# Q6 f) o; Q/ n4 j    for i in range(num_rows):
" j2 P. s3 Y$ u  c: S: i        for j in range(num_cols):
            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()
+ q# l; f2 l2 G    return axes

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

# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
+ M& Z7 w* x- _, t. o# {    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
" Z* g# {. H( B4 j# x9 v, E3 z
. J# \9 r' W) ^def show_images(imgs, num_rows, num_cols, scale=2):
$ _/ x6 F+ X) _    # 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):
0 b9 O) f4 [% h# Z. ~        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
0 d6 M8 u2 y9 R7 ~* x2 z2 v            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
( L% W6 ]* s/ |7 t8 w3 e$ |/ Y画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
4 C4 _5 K& h0 q- t% _
# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
" z7 U1 I& u: j: l- Q, Z$ n: vimgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
% `' m) z* H8 ^4 t7 r
- D2 P$ }0 h, U2 y- \
) m3 K5 p, E. n( K: v! ]7 n9 U
列出标签中每个RGB颜色的值及其标注的类别。
' K7 e9 D) }+ u* e. J8 t# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
# R8 q! R$ f5 v; w8 M; ]: n                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
* |* |6 `: s3 g- b                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
" E2 c' c$ l1 M. f                [0, 64, 128]]
# 标签其标注的类别
. Y9 z5 s7 t0 @+ x* B1 H, TVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
" O* o- }  W: `6 U, |' O               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
7 z; N% A7 e: x1 A  w3 g1 R有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

! o1 x: k! o2 R% Y" [, _有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
& t4 E) ]  p8 ?4 g- @5 D! o- kfor i, colormap in enumerate(VOC_COLORMAP):
: {: o3 i9 D. e6 x& t9 Q    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
* v6 b8 s' W& U4 K/ c    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

# 构造标签矩阵
9 n" O* w: h4 Y# {% B0 E2 ~. mdef voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

% z( P6 ]# j8 v& U9 o/ }可以打印一下结果
# }) _5 H1 l1 Y: l& l
8 W+ \7 Q1 r( L- b, h2 o; O2 g# py = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
" I! ^! X* u* t
0 k) p6 s# c' g! |1 X2.4 预处理数据

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

8 p8 B, d' f9 u  P; ndef voc_rand_crop(feature, label, height, width):
    """
    随机裁剪feature(PIL image) 和 label(PIL image).
; L7 j, s% [7 x$ y4 o    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
6 `* T: e* e8 |; ?5 i+ f        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
% g' ~+ \6 |/ v$ O( w    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
3 ^' X6 @5 C" A# W    """
    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
1 P9 ?8 p6 t3 d, K    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label
& C) e5 }+ [1 L- n" |7 Y" M
3 |6 U9 P* e1 m% i  r# 显示n张随机裁剪的图像和标签，前面的n是5
: ?( C# S5 r2 P8 `0 Gimgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
+ B& ?) c5 g1 m# |show_images(imgs[::2] + imgs[1::2], 2, n);
' t( ]8 t9 I4 u" F
  K0 x5 h, q- T5 ~$ E3 i
2 N" V. k# `! A% D5 j4 Y# \" B

  w5 C* M: ?; i7 ^

% a: F7 |" W: c- S' f3 自定义数据集类

3.1 数据集类

torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

$ \0 ?* V' s9 f__len__ 实现 len(dataset) 返还数据集的尺寸。
0 d2 i+ N& L, \0 m__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
, d' }$ n1 u! h" k  _" Y+ w
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
, s0 G2 r' i8 A" _5 l        crop_size: (h, w)
* t( I- u& S0 ~& C# C        """
2 `, }/ N9 T" H/ R7 {% D2 r        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
! U! G8 I2 i! U' q        self.rgb_std = np.array([0.229, 0.224, 0.225])
& k* t4 w, X9 e& i9 P, m        self.tsf = torchvision.transforms.Compose([
  x7 C5 @# W' A/ ]            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)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
. ~- j6 d0 R+ A0 W( i, B        self.features = self.filter(features) # PIL image
! R7 y) X  ~. R+ w/ a. V        self.labels = self.filter(labels)     # PIL image
" l. e3 [) I1 s! U8 a$ X+ X        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
- f2 U# i0 n7 w3 Z3 X
    def filter(self, imgs):
, @3 @0 M& l* ]- v. X# q        return [img for img in imgs if (
- p" p7 h  Q4 l" f8 y0 A            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
$ [' C/ A2 K$ I8 J  g5 ]
8 `7 S* h) b6 s) C3 c0 L    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
, Q4 d% [8 V5 e4 r; k( H                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

; Y+ c1 D& N1 J$ @' t7 d' |* [    def __len__(self):
        return len(self.features)
# W" M5 i6 v% F7 _3 p% z- D% f1 @1 G3.2 读取数据集

" _* f6 L4 t- p( {通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

3 p* n/ O% M. a# P2 B' D6 M1 Ybatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
/ J- n! }1 y) i& l% y4 ccrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
( [+ @+ r8 A. L. Xmax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
8 V" I: `. z! U6 e$ S" _
# 创建训练集和测试集的实例
- T5 O) p) E% a: u6 \voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
5 l6 t! _! L/ E7 W6 dvoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
& B) j3 l! e( I: W, `3 T& wnum_workers = 0 if sys.platform.startswith('win32') else 4
% T- c9 O# R2 _" B3 z. dtrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
2 P  V: g  n. K+ @- A, c6 H# s                              drop_last=True, num_workers=num_workers)
3 R4 u7 z- b' K" ^) M- ytest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
3 G3 \9 ~2 x" d                             num_workers=num_workers)
2 ?' ]/ @: Y  q
" R9 J  s6 i& q! w- c- v; f. y# 方便封装，把训练集和验证集保存在dict里
* G* ^# r5 [( S' d+ p! Ydataloaders = {'train':train_iter, 'val':test_iter}
) n/ A( C( z5 H1 x8 Gdataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

9 P/ e) P9 ^( |) i, b4 构造模型4.1 预训练模型

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

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

, Q$ N! S$ w8 a$ `6 U/ k/ Enum_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
5 O4 f, f8 {2 `' K! L8 P1 ?$ ?: o3 U
# 特征提取器
! t9 f9 x1 z2 u7 W' Z- T6 pfor param in model_ft.parameters():
! R# R# G: z0 M- U8 e5 N+ E# i" Q    param.requires_grad = False
6 z- K, O' S1 E; o8 d) _1 l, `' C! F4.2 修改成FCN
: o- E/ y! U: d; k
$ S& Z8 f) ?' C- K+ c全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。

对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
7 {$ G4 R( e8 \
可以先打印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 倍即可。

model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
9 ~) {) K* E! f& P% ~1 q7 J7 G              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) # 转置卷积层使图像变为输入图像的大小
3 p" B0 ]! A* O1 P  H" `' `
: W9 ^! r% R. N1 W2 ?7 k* }' [# 对model_ft做一个测试
! F$ T( V# Q! `x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
; s  Q% {* o9 x  C! [# E5 e# for X, Y in train_iter:
+ X0 x$ j0 {  a4 [- B% f#     print(X.dtype, X.shape)
( P8 Y1 o. Y& Y/ s+ v0 `" D. B#     print(Y.dtype, Y.shape)
#     break


4.3 初始化转置卷积层

在图像处理中，我们有时需要将图像放⼤，即上采样（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函数构造的卷积核的转置卷积层来实现。
4 X& F# C  w# a) K9 F$ @, y" y+ x5 y

# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
% p! W: ?8 k1 J6 o2 \    if kernel_size%2 == 1:
% S) M; g$ l( `. b        center = factor-1
" ~+ n8 {7 h2 J    else:
2 E/ s7 o; M7 f4 E, z, ~. q        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
! t2 d$ A8 ~! ^: G5 _0 H+ `    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
% S3 Z9 X* ]% s7 P0 S$ s& i. ?- i$ D    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight


在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于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)
6 E0 g. U( C$ w1 f* p3 L- i5 J


# c+ R, m6 t2 J% M1 d8 T5 S5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
$ |4 ]! `/ l; u) ]& G( E    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
* V0 r2 Y/ [; o$ I! C5 F    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
        for phase in ['train', 'val']:
5 i; Z# n. J( m, C5 F            if phase == 'train':
  B( c+ f$ |4 o$ H" g2 N- U% L                scheduler.step()
                model.train()
            else:
+ |( v! K4 g# ?. C. X                model.eval()
            runing_loss = 0.0
; m- q/ l, f6 ]  _, [9 s            runing_corrects = 0.0
* o3 l. M* u7 V- a. B" o, c            # 迭代一个epoch
* j6 \" p! [7 Y' B            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
                                # 前向，只在训练时跟踪参数
! [5 J+ {7 |! u                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
$ C; H4 l0 W5 u& b                    loss = criteon(logits, labels.long())
/ j; ^& W) N7 a. j# @/ Z                    # 后向，只在训练阶段进行优化
' p& J# ~4 H* n6 j' n( n                    if phase=='train':
                        loss.backward()
8 I. _5 q. H3 I                        optimizer.step()
) c7 t8 ^: x! z, _9 D: P                                # 统计loss和correct
6 ^$ a4 _3 T4 D                runing_loss += loss.item()*inputs.size(0)
$ a' T( d) c6 K& \( T                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

1 a5 B9 h# A, I8 I$ ~5 _            epoch_loss = runing_loss / dataset_sizes[phase]
  ?6 Y  c* G& N0 @            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
9 y! U% _3 j" u            if phase=='val' and epoch_acc>best_acc:
8 ~5 v0 v3 p1 F. M' G                best_acc = epoch_acc
/ y9 w! W! ^4 k, b                best_model_wts = copy.deepcopy(model.state_dict())
        print()
    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    # 加载最佳模型权重
, [3 Q' {. K: c0 v& y/ q7 J    model.load_state_dict(best_model_wts)
# q) l. m- p4 H. {    return model

+ e3 Y' a+ p& o' p! b下面定义train_model要用到的参数，开始训练
" P2 [6 r  V+ G  l: a' S6 O$ @
epochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
8 n6 d8 K* w4 goptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

+ V8 z; c4 p5 k# 开始训练
/ ~0 `6 `- G# _$ j' x5 b; w8 D1 H! r+ \8 kmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

6 测试模型

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

def label2image(pred):
+ e0 R- C: n- d5 o( ]5 [8 Z, d& B2 F    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
7 |" y% A3 Q, P& D8 @    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()

0 Z0 J, \/ y' B$ s

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

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)
; V; v3 [; }7 w) O4 Q$ E0 Jdef visualize_model(model:nn.Module, num_images=4):
* X/ G" l2 o5 e9 I  y( v* g    was_training = model.training
    model.eval()
    images_so_far = 0
    n, imgs = num_images, []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
            outputs = model(inputs)
            pred = torch.argmax(outputs, dim=1) # [b,320,480]
4 D* b* a+ a* U) a; V            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
. ?& E0 f4 t( _
            for j in range(num_images):
                images_so_far += 1
2 z6 _) [/ a' R                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
# J7 Z7 G0 X) [) z                if images_so_far == num_images:
                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
# g1 ?% k% h& e& M                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
% w7 E: H$ ?2 W  i1 `; _                    return model.train(mode=was_training)

# 开始验证
$ @: o7 T2 y( N3 @! hvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
& h+ J2 L/ S6 B" c* ~1 ~- P: [def predict(img, model):
    tsf = transforms.Compose([
! `0 x. E. z. r1 t8 Y# y3 T            transforms.ToTensor(), # 好像会自动转换channel
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
6 i0 {- Y' H) _9 T    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)

1 h! o6 K9 ?( Gdef evaluate(model:nn.Module):
    model.eval()
# [' d( |, X( r! e' D6 q* ?    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
1 c- c0 h7 A" d- s    n, imgs = 4, []
/ ?5 t7 P) F" W% B* Q    for i in range(n):
) f4 G1 l% P/ g' k) y, X4 B; v        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
, |2 e9 K; ]$ i+ Y( I5 J    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
9 G; s, ]# ?5 }: l6 E
# 开始测试
evaluate(model_ft)
" z; f) a# B4 |" B) K! ?
7 结语

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

Epoch 0/2
----------
, |1 X7 i. M+ G& t  v% V" a$ Dtrain Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456

Epoch 1/2
- s6 n9 S0 D" F. i; N) A* H* p----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

# X* w  u' H8 o; iEpoch 2/2
----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948

Training complete in 6m 37s

; B3 X. I6 J: D/ P5 H


: ], E+ c& k( S" F+ v当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
; C* a; u9 {6 V( E0 u! m
  W; ^! t% L7 n4 L% b. _( L$ g对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
' R! d; j  A1 X# i/ Y
" h$ o; ~, N, t: K. z4 B/ R
语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
  D+ P! }( i! S( o; O* e
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
2 }5 R4 N' A# t# a% j" L) YGCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
7 j4 B2 v% n1 _+ i7 HU-Net 用于生物医学图像分割的卷积网络
0 s$ Y/ y" n1 x8 K/ p1 B3 wSegNet 用于图像分段的深度卷积编码器-解码器架构。
6 ?! c9 Z0 R) H' r: w- k还有（DUC，HDC）、PSPNet等。

% t' A; L7 M1 Z( C6 c+ I常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

对于损失函数，除了交叉熵误差，也可以用这些：

Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
- ?& E* r' ^9 n
% K; u$ k. h% k
- _/ Q& b2 f5 b4 ^& x


————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507