Pytorch实战语义分割（VOC2012）

杨利霞

1 d6 j0 p/ o9 N0 g( TPytorch实战语义分割（VOC2012）
7 u) o$ u% X) K本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
+ J# ?. T/ v  S, O( B+ _/ C, L

! L/ O3 U, X6 G  E2 w% t! v语义分割中图像有关狗、猫和背景的标签
: T$ c. |# m- s" S8 _- D8 a文章目录

1 t5 {/ o7 ]2 U+ D: U4 _. p( t5 r1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
  h! U) Z  Z" ]7 Y, M7 Z/ z7 A2.1 导入模块
2.2 下载数据集
% ]8 u$ w! _8 j- U2.3 可视化数据
( D, m/ t- {3 ~3 s2.4 预处理数据
# |" I* a1 v+ i/ N. [3 自定义数据集类
4 j9 _! Q' x  e1 d' X  d2 ~, M3.1 数据集类
& f  ?) }; b" v9 G/ `0 g. I3 Q0 k3.2 读取数据集
+ V, q  }' S6 P  ]( k& Q$ b2 E4 构造模型
4.1 预训练模型
  q& f: L) C2 @8 w9 o4.2 修改成FCN
4.3 初始化转置卷积层
4 ]  M$ A7 g' B! S3 u: v# ?5 训练模型
6 测试模型
& _2 ~% \3 D& k* h2 a2 L3 E! E  `6.1 通用型
6.2 不通用
" ?2 ^+ E' a) Z, v) u5 x2 E7 结语
$ G7 b3 P6 [, ?0 l5 t- c1 图像分割和实例分割
* X' W" O* S% F/ O4 X3 [
- E8 |  D- f8 O) d3 Y) `计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
/ i# q& |0 ]* F6 J9 Q0 d8 f
图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

2 Pascal VOC2012语义分割数据集

# y9 f) L+ r# m8 x6 C' Z2.1 导入模块
import time
import copy
import torch
from torch import optim, nn
* @: G& b* i9 I: T1 x6 v5 @import torch.nn.functional as F
; `- X6 ~  d: Q; h6 oimport torchvision
( _0 C( e+ b( B4 O+ |from torchvision import transforms
( u( H" y, F9 u' }5 U( v0 mfrom torchvision.models import resnet18
% W$ O6 ]4 H6 o4 E8 Himport numpy as np
% Z5 L+ x2 B  hfrom matplotlib import pyplot as plt
& e$ m% D) d2 nfrom PIL import Image
import sys
/ q/ @5 V( l+ a5 S- asys.path.append("..")
from IPython import display
from tqdm import tqdm
2 w% Y, b" k& J- R: x1 |import warnings
warnings.filterwarnings("ignore")
/ v( J; N* N' p. p. ^% h3 E, L+ r
2.2 下载数据集

+ t% p- f: R. `6 ~6 d3 M  T% A6 D语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
% r3 i+ n3 o# O* j  L$ a- ?

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

定义read_voc_images函数将输入图像和标签读进内存。
! ^0 ]6 k4 ]% ^( f5 F& q1 C
4 _& T- R2 W, b+ O! Vdef 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
6 m7 e7 p9 s- t4 D) ^    if max_num is not None:
9 l9 g' E4 g' q/ d4 |! S        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
& p, [( k2 W7 x; ~! y        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
) Q/ c" D8 c- ~$ c5 L1 U7 b
0 K0 Q, Q+ O, }. j- t5 F定义可视化数据集的函数show_images

# 这个函数可以不需要
5 _$ J' e; c. l7 N) c% Ndef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
! R( w( `6 r; `8 ?  e8 O. }) e    display.set_matplotlib_formats('svg')
0 J+ m. H" z3 d1 d    # 设置图的尺寸
    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)
+ e" s7 [7 T! X    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
1 U0 j5 I1 O: I; ~% b    for i in range(num_rows):
* F* e9 Z" P( I7 i        for j in range(num_cols):
" C0 s. B; D' e4 f6 T8 X            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
3 Y, P' A) O' c2 v            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
: `& p) b: b8 r4 j    return axes
; q+ o. x% G2 C; K4 f0 [& q
! G4 a. C, R  v' A; o% T  I, U定义可视化数据集的函数show_images。
3 m$ `. H0 ]' p0 @( ], ^
# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
* i7 L6 P( w7 h% G& c. y/ }    # 设置图的尺寸
    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)
9 r7 |$ A. _. Q9 |- {$ @: F    for i in range(num_rows):
8 W4 D6 [" Z! [/ V# Y        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
0 v9 ~" A# ?- b, P            axes[j].axes.get_xaxis().set_visible(False)
0 ~8 R8 {: Q4 A            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

7 l8 B  j: ^$ K4 f) q# 根据自己存放数据集的路径修改voc_dir
! A1 f" [. H2 wvoc_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
7 _* W! X. W6 X) K* N- |show_images(imgs, 2, n)
. H( o0 x1 t( M) `6 O$ ]# k4 M

3 ^  p- I* B6 R5 |
列出标签中每个RGB颜色的值及其标注的类别。
* Y% T  f+ O( e4 \& V# 标签中每个RGB颜色的值
$ s+ h' U; |3 n4 m% Z4 tVOC_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],
. ?3 u" w& }8 t7 B+ H( P' \! T                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
5 V. z" \2 h: O1 v4 K9 G, t9 t; ]: g                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
& ?% G6 u3 B1 w7 }, f                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
# [! x2 h6 ^- ^# R               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
  ~5 g# V- n- M4 k1 [- s0 P9 D4 I有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
( p! [4 J3 w7 |
1 D: I$ L( I" s+ K有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
) T' e# i+ j: N) M; lcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
+ ]* ~( L9 x! jfor i, colormap in enumerate(VOC_COLORMAP):
" u/ i1 h9 a1 K) G7 _5 L% `    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
& _! W: r0 E$ V5 r" Y9 {4 b5 H
* F  ]7 q7 h5 O) M) [# ^7 [4 O# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
' g. w+ _6 e( B( l    colormap = np.array(colormap.convert("RGB")).astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

可以打印一下结果
* m3 n7 j  X. n: G, \" y
y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

' f/ @+ s5 ^8 L. N  N2.4 预处理数据

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

' F3 W! z' x: w, A6 \6 n6 J; D
% i7 N4 u) |' t, ~. Hdef voc_rand_crop(feature, label, height, width):
/ |, p0 B  W% |2 V    """
5 z8 A; Q  Y% r    随机裁剪feature(PIL image) 和 label(PIL image).
  L$ {. U0 [& `5 X7 j8 l, m    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
7 B" A, a; k6 \: ^    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
* d+ x" q# x9 W    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)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
/ G4 v. z) B) H  B. l' g5 W' y    return feature, label
) D! O* s9 b0 `
# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
' m7 B( |4 ]5 q& G. N4 n2 e* vfor _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);

2 t( s9 ]) f: O. i% R



3 L5 {4 i: U( O
3 自定义数据集类
# b! }/ b- |' [4 k* n* A( b# y
3.1 数据集类
1 v2 [% w( m* `4 d( D# ~9 v
2 K, }$ j# F7 @! ?* j0 Gtorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
, Q! k! t/ d1 c7 _# V- X" W
__len__ 实现 len(dataset) 返还数据集的尺寸。
; I- x" O& T3 K" E4 f% G5 S__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
7 [9 X5 q  s/ o5 z7 j/ X
class VOCSegDataset(torch.utils.data.Dataset):
0 ?2 `, o1 b0 d    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
8 Z! u- K9 V. g5 q0 K/ v9 L        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
& z  ?% g, D4 z        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        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)])
3 F3 u, _; Z. V, f9 ]' L) b6 C        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
7 {# S6 n" |: u        self.features = self.filter(features) # PIL image
0 i' w& k  A3 P: k) U* l" v) C        self.labels = self.filter(labels)     # PIL image
        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

    def filter(self, imgs):
/ a7 z% O% q+ p$ Q1 Y. l        return [img for img in imgs if (
$ N# R  Z9 n, e            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

    def __getitem__(self, idx):
  i3 A; v* _2 J, s* h        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
8 ~) x( g' V+ t8 s" d; D9 a( S' K                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

    def __len__(self):
        return len(self.features)
! T% c. `& n$ y; e5 a* {3.2 读取数据集

- @3 C7 u' `9 G) }4 s6 n通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
! c( V7 I/ D$ X" K
batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
$ Q3 }+ R* _  R$ Y; @- \crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
9 ^. N7 g; E: b$ lmax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
$ T- t# y: [3 h
+ o! Y; w3 s* I% N" Z/ W# 创建训练集和测试集的实例
0 H' K+ w' r+ g: yvoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
; y# e9 @  I( c6 }voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
' C% m& h7 H& ?% r0 V/ @2 C
# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
2 \% \+ ^) ]/ ]0 k  k6 `train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
1 X2 X8 ~- K* B0 {; o/ }                              drop_last=True, num_workers=num_workers)
9 W! [9 t8 U* G! c$ G1 `) `* stest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
( |1 W% b% k0 \2 b                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

0 v' z! w# X& A) R! u4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
; u. t0 R" t  C  {; w
! q( K4 A; `' b! N& Pnum_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
+ F6 m+ Y) W8 N0 y8 j+ xfor param in model_ft.parameters():
9 J, D: v3 v" B6 i8 r6 ~    param.requires_grad = False
+ J# q  f% X, G4.2 修改成FCN

" ]- ]: `$ {# H全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
0 E2 M$ x: o- Y2 [最终输出的通道包含了该空间位置像素的类别预测。
7 ?5 z! A3 [8 T% v0 L. W
- q2 l  \$ @; y对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
2 Z/ T( z1 Z) l8 k3 v9 S2 v! D
, Q! D7 s. H. n  W4 q# i5 o! A; [可以先打印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 倍即可。
! i& p0 H( `7 g, U0 m
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做一个测试
: ]  M6 D& p2 T& Rx = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

. v: \3 P1 X2 f% c) w* T/ i: c# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
6 f  F( z7 e* N* F" C+ h9 j% a# for X, Y in train_iter:
5 z4 r7 R" N( G/ U#     print(X.dtype, X.shape)
8 {: H! N# |# E9 `#     print(Y.dtype, Y.shape)
/ t5 q1 P( R6 u4 S#     break

2 z4 x7 }, }* C' G
4.3 初始化转置卷积层

* D; v4 e" x/ P在图像处理中，我们有时需要将图像放⼤，即上采样（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 }: \5 l$ G1 g! P& a8 R) o, |, M9 h
! P& k( B' [( X) M, i: b0 d+ Z
1 N4 Y( n! [7 `' S2 U! P+ l# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
; ]" i0 K: e4 D! U    if kernel_size%2 == 1:
        center = factor-1
    else:
        center = factor-0.5
/ j, A5 W: w& O/ p7 z    og = np.ogrid[:kernel_size, :kernel_size]
5 W9 E$ q. M6 ]6 x7 n" ]2 n    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
! L3 o8 H( C6 }5 K    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)
    weight.requires_grad = True
    return weight
) T  K( a& e) i

在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
6 b) M, _" y, j2 H9 a, V8 Dmodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
4 e. ~) C) _6 p% D% B1 R" i6 C
+ R% p- y5 i  m  c+ h
! B/ b: I/ E+ ~, [/ k. Q
5 训练模型

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

5 _9 S3 @* Y. I- [1 Qdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
! s7 Z3 a3 Y* ~/ t+ O8 [    best_model_wts = copy.deepcopy(model.state_dict())
4 `# y! k- R; p# v( e2 s    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
/ h/ H" a7 {: j9 S/ A        print('-'*10)
- {; \8 N6 d2 h7 U6 S4 Y        for phase in ['train', 'val']:
9 A+ ~& ~( f* k4 ]4 S            if phase == 'train':
  {$ }2 C" `! g4 R4 `7 R1 b                scheduler.step()
5 h8 A: C- s& V                model.train()
/ [3 y1 A7 D' [- B6 w: S' Z4 S            else:
                model.eval()
3 R; J" d0 B! }0 H            runing_loss = 0.0
            runing_corrects = 0.0
" n4 N3 S, s3 e' |            # 迭代一个epoch
* E+ G1 o. {, {  ]$ l/ b            for inputs, labels in dataloaders[phase]:
% u  }. h& s) b; O" G6 K                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
* @" e, ]3 I( A$ {* `  {                                # 前向，只在训练时跟踪参数
: y) Y3 p" J) R" u                with torch.set_grad_enabled(phase=='train'):
* D. |( B! g5 @9 G$ {7 ~                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
) P5 c5 T$ Q5 H2 T                    if phase=='train':
; c$ U) L% M! h! e9 n0 T* h. k                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
0 c' e: g7 F% G4 _' Y0 O' g3 K
            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
- X  ]) P4 X; E            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
' r! j6 |' r' {9 i' e                best_model_wts = copy.deepcopy(model.state_dict())
        print()
' K3 H! a0 _3 F& j# c7 E* R( f    time_elapsed = time.time() - since;
9 h- k  k9 T# e( h3 q! _  N    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
1 G8 h) H# Z/ R6 J    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model
' b3 K8 l" k$ f9 Q" U
4 W2 c. H- v4 L) v$ Z下面定义train_model要用到的参数，开始训练

# V7 w4 }/ m- D( Z% i: oepochs = 5 # 训练5个epoch
: \# R) {( @8 V0 Q- P$ Ccriteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
9 m1 @3 [& N8 A. i: rexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 开始训练
4 G( I( A) y: L! d! ]+ j- {  c6 Lmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
; d/ b) [. n/ R, }6 {
6 测试模型

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

def label2image(pred):
2 C/ g- P( J4 d8 O4 N; ]5 ?    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
; A5 `  V' U4 t1 m- f    return (colormap[x,:]).data.cpu().numpy()

# j0 \! X/ a* Z& S$ _0 ^

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

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)
def visualize_model(model:nn.Module, num_images=4):
2 j$ t6 u% [! ^3 M) [    was_training = model.training
    model.eval()
" m# f4 |7 U  u    images_so_far = 0
- r4 E: N; J. {2 j: x3 e( A7 w7 |    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)
" C; b, A+ T: `3 o# H% Z) w            pred = torch.argmax(outputs, dim=1) # [b,320,480]
$ [. N5 Z, \9 g; [            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
3 o) n' t# ?! [) p
            for j in range(num_images):
: e" Z5 [- ~' G0 G$ ]& |% \                images_so_far += 1
* B4 o+ I% Q. ?7 g! a                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
+ Y# ^6 u" \7 Q& F$ t                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)
4 V8 z! \" y: f" v* T, D7 H% \
1 y$ h' ~2 }0 s7 P) ]6 G# 开始验证
0 {/ I# a6 c9 V. G7 h/ r9 Y! z- Hvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
$ r4 A7 ?3 I" p. K' ?1 b    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
2 `6 a6 Q0 q$ P$ v            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
: Z7 a6 S) {7 \0 z' \3 v    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)
% {+ V4 i( i6 s- O- J, B2 s+ x, d
, Q) t9 ]( f2 j/ G2 ]3 W7 T% W1 Edef evaluate(model:nn.Module):
, E6 o: Y3 o! w    model.eval()
, d. v  s5 B, V- x$ H# w* X    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
+ F; x3 G' W+ g+ Z) u( Q. }; I5 V    n, imgs = 4, []
    for i in range(n):
        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
. q) K8 F* C% N        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
( d5 Y  L- y& r6 w4 D    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
. u) D# s! T; ^* M5 E
5 S7 L8 Q7 x+ t8 g3 y: X# 开始测试
evaluate(model_ft)

7 结语

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

Epoch 0/2
& o2 ]" H' j- O7 ]: m----------
; {" q  T- J3 G7 \! b/ N' Vtrain Loss: 1.7844 Acc: 0.5835
6 S! g/ Z& M* J" q* q+ Nval Loss: 1.1669 Acc: 0.6456

Epoch 1/2
% e# J. D: X# s  f----------
5 H- q: M4 p8 o* A' _. ?5 Ptrain Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929
+ U0 Z) w8 G" X: w  n
# R- x, ^' t" \0 P. q/ D/ _Epoch 2/2
----------
2 H8 G+ g, |" B& H! ~train Loss: 0.9578 Acc: 0.6706
8 x, z( W$ f2 N* L! M/ H8 lval Loss: 0.8088 Acc: 0.6948

Training complete in 6m 37s



) m* C! A, K9 b" K
7 ]' V2 f2 H" o- b. k0 _0 ?当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

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

1 C; F& A; B4 `  G% H
9 B' |( Y& G' b2 }0 k. I* N! s/ g语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
3 |' f$ W0 x& B2 B* Z: v0 F
1 |' j" B' H# g' R2 _Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
  v/ D- ^9 M' j  {- S/ XGCN 通过全局卷积网络改进语义分割[论文]
( T6 b5 b8 A2 I- [: }( o/ FUperNet 统一感知解析
9 |  s; @% J; @7 r7 W% y7 kENet 用于实时语义分割的深度神经网络体系结构[论文]
, s+ `* ^$ K/ [6 ]. \) PU-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。
0 j6 e+ K' P* h; o6 n0 ]
常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

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

Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
' v7 m# B! Y, q' hCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
) r; u8 J0 f* t$ d1 F3 ^Lovasz Softmax 查看论文：Lovasz - softmax损失。
, ?7 L' U4 |1 h

6 O& s- \2 }4 W3 O

2 T- t9 [; c" o& B9 |' ^: t
————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
/ Y3 B: x7 W. t6 s

9 C$ i" R* c$ l  E