Pytorch实战语义分割（VOC2012）

杨利霞

1 w3 s6 E& {( e" P# K
Pytorch实战语义分割（VOC2012）
+ ~# z) e+ h1 a+ `5 Q本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
% R# O0 P7 `8 ]; L) F语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。


! H) }. J/ C1 ^8 u4 m语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
  R5 \2 G' h0 p, m$ x2 Pascal VOC2012语义分割数据集
  F! M! O4 R* B' o0 S2 r# _' H2.1 导入模块
2.2 下载数据集
2.3 可视化数据
: G8 t; t# c( {2.4 预处理数据
3 自定义数据集类
3.1 数据集类
2 T# I+ t  {9 V0 m" U8 g3.2 读取数据集
4 构造模型
+ i) k- ~3 `5 I' O6 l4.1 预训练模型
% u3 V5 r' K% h$ k. x& C, j$ Q% F4.2 修改成FCN
2 r1 o. b& o! q2 X" _4.3 初始化转置卷积层
; c2 `: ~" ^& e3 j; r( s  H6 V3 h) g5 训练模型
7 Y& s9 D. H% v2 x6 测试模型
6.1 通用型
6.2 不通用
; G! _9 m* b9 h7 结语
1 图像分割和实例分割

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

- d9 L7 F+ ^% c  Q( D; ]2 Pascal VOC2012语义分割数据集
2 p1 o4 G: I! O5 _
2.1 导入模块
import time
% x7 _4 j( ]& g& N; h% p0 l1 [" Jimport copy
* c' b% A" A# O5 r) nimport torch
, T' ]. s5 ^- W" I6 F0 E6 K- Hfrom torch import optim, nn
import torch.nn.functional as F
* M% }  O( c* N- Kimport torchvision
& D+ p0 S" T  o8 b) k. w( Wfrom torchvision import transforms
from torchvision.models import resnet18
  Y9 e( B/ \$ F. H$ iimport numpy as np
5 S0 y' Y. E2 I. s# I. Z. sfrom matplotlib import pyplot as plt
7 j" R8 B( b5 i% lfrom PIL import Image
import sys
+ `: W) Y7 G9 d. u' Tsys.path.append("..")
; G0 J2 X% o$ ?" u8 T5 A5 k; {$ yfrom IPython import display
2 Z2 \  M. }1 |1 `' _, r! P4 c) wfrom tqdm import tqdm
+ n: k$ f; s& S& u% Q: Bimport warnings
! b2 R7 T) b9 J; p; j& R' Ywarnings.filterwarnings("ignore")
# W. V4 A/ D% [8 x3 G2 c
2.2 下载数据集
% E' t8 R; T8 O
语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：


ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
. h4 m; v6 o' t" ?2.3 可视化数据
  M) p& m: y5 M8 W
6 t4 A, S% ?2 A- O& E定义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')
0 n( f; v" {* p    with open(txt_fname, 'r') as f:
        images = f.read().split() # 拆分成一个个名字组成list
    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")
; ^& [* ^6 l2 W4 R7 l        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
/ S5 o* H) j* c0 B. W+ w    return features, labels # PIL image 0-255

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

# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
6 b. a$ T( k5 I5 }5 p& R    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
& @" A) J/ k+ _; q4 D    # 设置图的尺寸
/ ?" M# |( J- S9 q) u, B* V: `    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)
) \2 w1 e) p  L8 p) o- Z    for i in range(num_rows):
6 M5 a' A# s3 M# O        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
, W0 I# U/ k1 f/ x4 J            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
) N  G0 k- w( m  Y    plt.show()
    return axes
; ^( L; h1 e% f$ e
定义可视化数据集的函数show_images。

# 这个函数可以不需要
- o% t; n+ M' ]! mdef set_figsize(figsize=(3.5, 2.5)):
+ @- N( R( B2 }- `2 \    """在jupyter使用svg显示"""
4 w, R' o9 Y+ z  g* _7 D- C    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
! ~0 b: ~+ n  R# ^    plt.rcParams['figure.figsize'] = figsize
- n, D% h& a3 _# I
' V. ^) ?( W2 S8 ?" odef show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
" d" I5 S; w7 T3 l9 _    figsize = (num_cols * scale, num_rows * scale)
- D  y' q/ \- W# k6 T* q    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
0 ?) P  |+ k* v    for i in range(num_rows):
        for j in range(num_cols):
! f' [# e' U0 K/ G% Q& ~            axes[j].imshow(imgs[i * num_cols + j])
6 ?  s$ f! E: g, k6 z            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
+ R, T- e1 S3 Y3 i% p    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
, E$ \) P8 |+ V  A9 D) q! F
! B# Y: A" ~3 h% Z) w# 根据自己存放数据集的路径修改voc_dir
5 N- D4 E7 P7 F+ ^voc_dir = r"[local]\VOCdevkit\VOC2012"
( i1 i9 X4 I- Btrain_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
7 W* U9 I% ?6 w, ^- V2 t4 f" Hshow_images(imgs, 2, n)

; n: y  y: D1 w) O8 j

列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
- u0 |. N' E0 k7 l4 WVOC_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],
. q& x$ t. H! a! p                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
# K. n: o: {3 A- q                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
7 Z9 ^; ?( p3 c  S# H8 Y# 标签其标注的类别
# n% C4 Y2 \, J. K$ wVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
; t9 j/ ?# N! l. A6 {+ j               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
" V" k; b. z8 U6 d9 o               'diningtable', 'dog', 'horse', 'motorbike', 'person',
& f: \- C5 t$ [  _5 l! J: c$ K) x               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
* D) L- _; ]* I, n有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
4 y; r( ^  {( `# _$ `' N1 Ycolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
for i, colormap in enumerate(VOC_COLORMAP):
. @( e0 B7 k! v2 N5 H    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
7 T, |. |& P$ u' O
+ ^# }! ~7 R- s0 K  p/ z' ]0 {# 构造标签矩阵
! e6 c. Z8 ~; v7 s1 t# S4 pdef voc_label_indices(colormap, colormap2label):
+ g, M  [5 R8 X! I, H    colormap = np.array(colormap.convert("RGB")).astype('int32')
" i, }3 r4 o8 _5 h, w& g! P& X    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

+ w2 {7 E; {2 Q4 k  G" `可以打印一下结果

y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据

在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
% v; H& D3 l! b8 s) N" G/ r! }/ F
  P3 z. c/ a) p2 D+ A7 K& e0 y
def voc_rand_crop(feature, label, height, width):
8 e1 O; Y! N5 J9 S# n& p. G7 X3 `    """
) @8 x) T% K* ]" E    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
4 p, @$ E/ f- F* W  A    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
1 \3 `2 t( C5 @, R' c0 w    label = torchvision.transforms.functional.crop(label, i, j, h, w)
6 g6 }; K% |9 i! O! C    return feature, label

5 l! d+ A3 h% G' O# L0 S# 显示n张随机裁剪的图像和标签，前面的n是5
( X5 _! ^  c* ~# E. \* L$ aimgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
% _, q' }) Z" u* Eshow_images(imgs[::2] + imgs[1::2], 2, n);
# G* \/ H- n7 w, [/ _, E( _


! x7 O- g+ B. }! G7 ]
( Q6 u3 m$ X8 n1 a# s" L1 H" P
2 n0 {& H7 n. X5 m% ^0 @
3 自定义数据集类
' O$ _2 Z) t- H' s3 d  `7 a$ f- c
; k- k" D7 G2 N5 B" g9 I$ h3.1 数据集类

torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
, j" Q' W( i8 d) C
) ]# C# e6 A; A! t* C6 I/ W__len__ 实现 len(dataset) 返还数据集的尺寸。
: N2 C# o: ~" ]" }# J' ?$ c% ]0 W__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
2 `& s- L2 Y+ H由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
  H# d4 m: @$ @) r4 a- u6 X; q$ O( o
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
+ J! Z: K- r3 v        """
6 Q; d+ ?' N; \) ?# z0 s7 r        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
* v$ [2 _) i9 B9 o2 [        self.rgb_std = np.array([0.229, 0.224, 0.225])
! z& z6 V4 b3 d: ~: d: t+ _7 H        self.tsf = torchvision.transforms.Compose([
; o; f3 _1 e" ?  G4 }$ f( v( A            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
; h: S: A  N" L$ W+ I        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
& `9 `& O% R: W, e6 v* R0 q' |. t        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

7 F' r' q3 M" J2 w4 Q" x+ z% d5 [    def filter(self, imgs):
7 y2 W, E1 s4 F/ u2 C2 K        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

& o0 {+ ^2 S; c( J# a    def __getitem__(self, idx):
* }3 y# t8 t- ?1 R5 @5 A        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
; b- q0 ~& `6 t+ w, ?: }                                # float32 tensor           uint8 tensor (b,h,w)
  l. \* O6 R2 W        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
4 _2 t. r" e+ E+ d& |* B- o* H
: J# }2 r, e8 b4 u$ u" G    def __len__(self):
! b2 g% K& d) U        return len(self.features)
5 L) }9 e4 f0 ]4 A( E" ~: D; L3.2 读取数据集
5 z8 \1 t4 P8 y* @; U: g
  ?' J7 R3 w0 m  ^0 ^! ]通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
0 D' v4 H. I- X# `; s/ `  ]max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

9 J4 x/ Y9 j; o9 I8 k) N: U) v. G0 X# 创建训练集和测试集的实例
, f- y* v$ M- Vvoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

- b% J. u% B3 E- i$ g# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
8 N; i$ D; R( d/ g                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
                             num_workers=num_workers)
0 B- @- _, e* E" W/ t; d' g7 [& L4 G
3 P7 j2 y; h5 S' l& {& k1 h2 h# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'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')
' l, ^) m8 r: ^5 q
  M1 N) ^  L2 X# p5 o9 vnum_classes = 21 # 21分类，1个背景，20个物体
& R& |, w1 k' o7 d) H) x/ h. Qmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
6 n/ a% k% x7 ~) e
# 特征提取器
for param in model_ft.parameters():
/ `# a4 z5 O* o0 q" L  U    param.requires_grad = False
4.2 修改成FCN

9 m! Y. Q5 @+ V- t' T/ z全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。
0 r# o, J5 G7 m; o1 m
8 y& P5 }" `" I  h' d" o/ T) e对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

可以先打印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 倍即可。

- m" I' ]! e5 c1 v! q# pmodel_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])
( d( n: S* u5 C) O) ~
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
2 C) g4 ~! z, T3 t. W#     break


4.3 初始化转置卷积层

在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
( f( F5 c- ^# b6 c在坐标 (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函数构造的卷积核的转置卷积层来实现。


0 t6 H3 i/ J, K& T/ r! L* W# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
        center = factor-1
    else:
5 l, [4 X* B# Y! r% G( N; L7 b( O        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
" X: E0 a7 d; r. b. A    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
$ Y0 N& v, {" w, Y# ~5 J  {, |    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
( d! I' i  s+ D    weight[range(in_channels), range(out_channels), :, :] = filt
4 Z) l6 v5 @' C* b( D! L    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight

2 f. z+ Z- G$ r+ ^# e4 t
在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
8 N. V- }2 ~" A9 i
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)
; g5 s( D* n/ _3 r( @
# h; f, @+ \! U$ d  F# \- i5 P& i
0 j' j) p4 k4 g2 W5 W, n
! g- g8 H4 e* m! [; Q7 N5 训练模型

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

- W- L% B% I9 l& ndef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
. X4 o) R4 L0 t4 [" ~+ X    for epoch in range(num_epochs):
8 X# [* Q* d% ]$ m/ \# |        print('Epoch {}/{}'.format(epoch, num_epochs-1))
% a- h1 J- b2 I: @- \; [- n" Y* G        print('-'*10)
        for phase in ['train', 'val']:
1 ~* \. A6 e8 R5 F            if phase == 'train':
& n& x# T! e/ E/ c8 a; w                scheduler.step()
* ^) j1 ~1 M5 {5 K# [" ~                model.train()
9 k. \" y; n+ m' O: i* e' P            else:
4 O5 Q% d2 e* E( F, v                model.eval()
            runing_loss = 0.0
            runing_corrects = 0.0
" W. a+ q9 k/ y8 g1 [* U1 ~5 x! t            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
; |5 S% v8 P0 a9 _                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
0 v0 ^# C8 Y+ p, ]8 n' t                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
  ~+ j/ T9 s7 J% M% h4 X/ t& h                    logits = model(inputs)  # [5, 21, 320, 480]
7 x1 U, l6 v, J7 W6 }2 i9 t" O                    loss = criteon(logits, labels.long())
* B  T2 w! r3 ^9 {0 s                    # 后向，只在训练阶段进行优化
- L9 G" s( P( Z4 t# G                    if phase=='train':
                        loss.backward()
) x. X  X; o" L* P1 [  ?                        optimizer.step()
                                # 统计loss和correct
$ K. B. ?+ U/ I9 I* U                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

8 V( X* p/ u- ~! F2 J            epoch_loss = runing_loss / dataset_sizes[phase]
, e5 v! O6 i) D) o            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
0 K/ l( |/ H# Y            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
        print()
, Z, V5 P1 u* S$ A) U# k9 d    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
' I2 x  V' u8 V( B6 `9 K    return model
/ s2 I5 A4 x# c5 ^/ v  w; U. n
下面定义train_model要用到的参数，开始训练

epochs = 5 # 训练5个epoch
: J0 r" p' W* _  I( \$ jcriteon = nn.CrossEntropyLoss()
/ l5 `  n: T7 B. g* qoptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
' W' ^5 U( g  |exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
' B* m; S0 y; T$ F' `+ f. J
# 开始训练
& v7 _. Y& Z. ]) D* w! O7 E8 Umodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
& }" s8 `) c5 k% |, R
6 }$ y7 p% l* u7 z6 P6 测试模型

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

def label2image(pred):
    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
6 \9 U$ ~$ v3 x; ]. m0 ]    x = pred.long()
: V# w* y# i5 a% K2 q    return (colormap[x,:]).data.cpu().numpy()

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

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)
3 }2 i" m% V8 O3 w# ndef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
    model.eval()
0 M- z9 d! P7 O    images_so_far = 0
; d" r& _/ F! w0 O    n, imgs = num_images, []
    with torch.no_grad():
5 z/ B6 X. G5 E        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]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
2 F  y  d/ p% N* M' l$ f6 Y
            for j in range(num_images):
$ x8 c) `9 ]6 {/ N* m4 Q% d                images_so_far += 1
                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
1 f8 P, s" W$ r& A* A6 u* Y                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)

6 s% v( L# }9 E# 开始验证
3 F; n; x# U( V* P9 s) Qvisualize_model(model_ft)
% [* J" r( p  Z5 o; Z  ~) D6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
7 w) s; O& }- F# |2 ?3 j! {5 Z" X' N    tsf = transforms.Compose([
' i3 S9 s" b. [/ |( Q1 [: S) T; l            transforms.ToTensor(), # 好像会自动转换channel
4 q; `( b- Y9 ^; [            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)
) D% k; o% z. }% E; F    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
: }: j# [6 ^9 h% `& _5 X
def evaluate(model:nn.Module):
$ C) E8 `! \- g! U  k    model.eval()
; T1 V) v, }' [1 u* {( b% y    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
! W6 @: v5 @) }; f    n, imgs = 4, []
    for i in range(n):
        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
- n, u' ^- q3 s! ]# E        pred = label2image(predict(xi, model))
* ~5 B9 f2 u& N) P' B        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

# 开始测试
1 Y$ P1 ~( r$ _( Cevaluate(model_ft)
/ I$ Q9 s, k% o' m/ G* w+ I
8 j3 \" |+ {2 [- p' Q( Y7 结语

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

Epoch 0/2
----------
train Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456

: D* A0 q( S7 [6 MEpoch 1/2
, k" a. y# D8 N4 V; _/ Q----------
' X9 f" ]& B8 Jtrain Loss: 1.1288 Acc: 0.6535
8 p8 Y* t4 U2 p! |' Uval Loss: 0.9012 Acc: 0.6929
: C" ^% q# K& }  L$ V5 Z
3 `' z; s& g; O  Z0 J, R0 U, y" kEpoch 2/2
----------
train Loss: 0.9578 Acc: 0.6706
& X( A8 H  t  B% r7 Q, v4 \val Loss: 0.8088 Acc: 0.6948
0 k" p' o  m3 B# J
Training complete in 6m 37s
  _& r* `! v# G1 g0 U
) O7 p) U" z4 \3 l
. I9 K* R' D2 A4 ^2 U) ]
2 v$ @" b0 x$ M6 n
2 z, x$ }+ J# B! _当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
" h+ w0 O$ n7 P# L5 C& Z
对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
0 r1 A) r$ H% }" A( ~
4 e, E4 A" s. `/ {4 U, E& o) U4 G
0 H/ Y: O* H( q2 _$ R, y语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
2 U/ g8 _1 d2 Y% D5 h8 n2 v- ^- kENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

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

对于损失函数，除了交叉熵误差，也可以用这些：
0 {. V4 K# W6 @  a
Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
: ?. G' q8 s8 `  `; o0 ]& mCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
9 ~: v. ^: p% XLovasz Softmax 查看论文：Lovasz - softmax损失。
1 u; o) M. s  G9 U" `, o& G# d$ v

  V$ B* @% w. K5 z# ^6 r4 h8 ^
5 |' i6 [. I% ]* ?

+ I$ X) A- N- z————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
1 f; V. Z; H& u  g) f( n" B
3 {/ _. G: [3 l; V3 d2 A
/ o6 O$ [* E7 v3 F% U