Pytorch实战语义分割（VOC2012）

杨利霞

2 A& L* V( D# \  W' a: b
5 @. H, q' v, o; e( T) L; r3 @" e: l; qPytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。


  s3 m5 c4 k% o! L' s语义分割中图像有关狗、猫和背景的标签
- C% [9 O3 V4 u- _/ \8 z- k0 K文章目录
0 y( o$ i+ s9 t( p, I: I
1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
8 b" K% E0 f0 Q; p6 F2.1 导入模块
2.2 下载数据集
8 f- ~. Z, D9 X+ {- _2.3 可视化数据
8 ^, |* h$ S3 _; T2.4 预处理数据
3 自定义数据集类
1 _2 T* ?" R8 |6 c' o; B3.1 数据集类
3.2 读取数据集
4 构造模型
4.1 预训练模型
/ n* Z: j# V. |7 y7 X: J1 _4 G7 \4.2 修改成FCN
5 ?3 W2 k4 C! J3 k1 [+ H3 x0 Y- F4.3 初始化转置卷积层
! _9 s4 b- n. e2 ?5 训练模型
6 测试模型
6.1 通用型
6.2 不通用
7 结语
1 图像分割和实例分割

- Z4 |: Y& U5 e1 Y: Q计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
) K% p# S! p4 o" _( H) f. @: G9 l$ j
# L+ y$ `* \# b. L图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
  K2 T, T  k5 a( K实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

' P* f7 X7 n+ O5 c4 m2 Pascal VOC2012语义分割数据集

2.1 导入模块
import time
/ q" |# r" i0 Kimport copy
import torch
/ E- Z) ~. {! j4 K" S3 Qfrom torch import optim, nn
3 Z2 B$ Q& k% }. j  Qimport torch.nn.functional as F
0 E( `, W+ h! p7 J- v$ oimport torchvision
7 `9 _& P* t% }( @, A& t0 Wfrom torchvision import transforms
from torchvision.models import resnet18
* q7 Y, C; y4 p: w) x3 l/ n5 uimport numpy as np
$ k- t5 i0 Q7 X. Nfrom matplotlib import pyplot as plt
from PIL import Image
6 {& T6 z( h4 C. T9 Uimport sys
/ [5 u! B0 @) Z# Y; E0 D! vsys.path.append("..")
# W$ ^$ c; K$ k1 R8 i3 ifrom IPython import display
% t9 r4 e0 k; I& L) ^3 Rfrom tqdm import tqdm
4 T$ E2 i$ p& D1 S7 Z5 Simport warnings
6 q1 Q/ ~, _& ?$ f# q2 h/ R) wwarnings.filterwarnings("ignore")

2.2 下载数据集

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


ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
2 Z- m  E4 _0 u7 }JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据
3 i2 P1 u1 c, ]
定义read_voc_images函数将输入图像和标签读进内存。

7 u) N0 u' y, J6 p9 }def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
7 a  `& }8 V) h: ^# \    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
    with open(txt_fname, 'r') as f:
- n5 a! T% x8 }8 V7 v+ w        images = f.read().split() # 拆分成一个个名字组成list
    if max_num is not None:
- |  E9 o1 P; q1 D5 j        images = images[:min(max_num, len(images))]
# {) q7 V0 d0 m    features, labels = [None] * len(images), [None] * len(images)
7 f( J/ V! D# v& f    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
% d9 x0 T8 ?) S/ C5 s        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
- S( K: O2 F0 c1 L, _        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255
* l9 Z* k8 Q' y
定义可视化数据集的函数show_images
, c9 r, h' w, @2 v: @- i
# 这个函数可以不需要
' h8 b8 ~$ F' y+ i) V$ h* f) p* e+ ydef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
8 M8 h! _/ r  \9 j- }    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
8 z  y2 x* b! R    # a_img = np.asarray(imgs)
) L" |# S/ }3 h7 x    figsize = (num_cols * scale, num_rows * scale)
! ]% E$ N1 F3 C3 a    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
( }" m9 V: m3 b9 E6 L/ \    for i in range(num_rows):
+ P" x7 B1 n; A* i" F        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
3 j' l7 C% D) Y% V* Q            axes[j].axes.get_xaxis().set_visible(False)
( E0 ?3 x- A: o; S/ a            axes[j].axes.get_yaxis().set_visible(False)
4 W2 k- r* M) f% b5 }    plt.show()
! ]* i$ o' x8 l1 l& G- ~$ f$ U; C; }    return axes
. r$ V9 q# @1 q4 n$ N
定义可视化数据集的函数show_images。

# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
8 \, B5 ]6 s9 E2 x5 L$ r    display.set_matplotlib_formats('svg')
/ T  e7 h" f4 u: h6 R    # 设置图的尺寸
# i2 P! P, J7 u/ s6 L    plt.rcParams['figure.figsize'] = figsize

' n+ p1 }! p% wdef 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)
    for i in range(num_rows):
( b" h. S; w( {8 a        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()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# 根据自己存放数据集的路径修改voc_dir
: E0 c! K8 {, H9 T% V% t4 ]- B: Fvoc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
6 d- M+ {0 e1 on = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
9 o! z, A7 n* k+ X& \; Xshow_images(imgs, 2, n)
: B: a  F" h! K: f
. H2 Q: h) L0 p2 \( E, z, w
, m3 p. j) C0 X. \# x
列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
8 r. s1 f& H* }$ W: _$ p( PVOC_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],
; S& g! O; O; b) D6 }# M( |                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
+ l* y/ \4 m6 d2 w: L# I                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
* g2 f6 B! a4 _                [0, 64, 128]]
. u& A8 d0 n: `! E& \, M  C# 标签其标注的类别
9 h# k8 T: j6 P& W  xVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
! T& D  x, M: }& D               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
  ]* F! b" n* x( H, E0 M               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
8 B' [4 o' `; u  ]- }/ Z$ ?
6 U6 D! w; Z+ m! g8 X6 V有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
5 P: s& G" C% `9 [- L; Mcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
for i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
, Y$ D; `; h9 B* R    colormap = np.array(colormap.convert("RGB")).astype('int32')
) S6 [" p% V( C4 j( J    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
1 J- O1 p- T- y/ g- k    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

. z  V# ?, ?" j& p可以打印一下结果

y = voc_label_indices(train_labels[0], colormap2label)
% B* O6 L! L6 U" w. \# Lprint(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据
! o+ C2 b  g; S' t1 A1 C" X
8 U4 Z3 M4 v* n4 V0 p在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。


def voc_rand_crop(feature, label, height, width):
, M' c' I& O. S5 A; ?/ t9 q    """
: r2 U& A/ h+ H4 n0 ^  F  n    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
4 z( f" |6 v/ ]    Args:
        img (PIL Image): Image to be cropped.
' h7 l+ P8 h$ }7 D        output_size (tuple): Expected output size of the crop.
    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
7 V) y+ G& h+ U2 V    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)
    return feature, label

# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
, D1 G9 d" M' t, h8 ~: ]; j! F/ |( M7 ~    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
% s5 T* k: M, w! w; `# A  I4 ^show_images(imgs[::2] + imgs[1::2], 2, n);

* p0 F4 M% Q% g+ Q; l# D




3 自定义数据集类
9 H1 C" h) [$ `5 S9 B
3.1 数据集类

6 g- E+ g7 U6 u% N  k3 e7 |  storch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
( A, p3 ^* m4 g" r1 p" c
6 H# G" X& w5 \0 R4 B! T, Dclass VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
: L7 g* h0 W  `* k+ q2 c        crop_size: (h, w)
, X$ q0 N# T, U5 j9 m/ b& r5 h        """
        # 对输入图像的RGB三个通道的值分别做标准化
7 B- a; J! v% i* k$ W3 R        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
  H: v5 r# `" b# z5 @3 X7 a        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
        self.crop_size = crop_size # (h, w)
" V  u+ q1 c/ B9 k1 P; u        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
$ m; A, A! V; z        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
9 ]; n3 s! ~* `$ |) N: J        self.colormap2label = colormap2label
! v+ r* m3 E% F$ c        print('read ' + str(len(self.features)) + ' valid examples')
/ H7 H$ g: J6 r7 U% B
( m4 Z/ Z! e% T  {5 e9 K    def filter(self, imgs):
% O, w/ s$ Z; `* p1 b        return [img for img in imgs if (
( f9 u% x8 S5 c8 P            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
# L% `5 }% P0 o0 t
4 V3 D  ?. n0 B# p) A; J* [    def __getitem__(self, idx):
6 p/ f. T. Z' w* @7 R+ d        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
                                # float32 tensor           uint8 tensor (b,h,w)
2 Z0 k* a6 y4 K- |5 c        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

    def __len__(self):
        return len(self.features)
1 j( y* |" m4 d4 t3.2 读取数据集
; b$ N3 Q+ G3 W: A. w- ^
( |( E( c. C7 K0 A通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
2 N* U- D( o  \2 W* p. L8 S7 N
batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
9 Y& s- c* \  R  W3 \# U$ R) ?; B
: B) n2 U" W1 o+ @; D6 m# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
! b3 H0 n- y) z+ R( n; ivoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
* U" I' i, r4 b; _/ ?  r
# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
- A$ [% M8 ^. Z: e3 v3 q1 mtrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
! i( [( t: w3 ?% ]! v: jtest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
                             num_workers=num_workers)
6 k+ k" I" Q4 Y5 Z
# 方便封装，把训练集和验证集保存在dict里
& f1 v8 V7 @) K+ ~% ydataloaders = {'train':train_iter, 'val':test_iter}
dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

# u+ Z9 z$ ?$ s" \( P7 ?7 v! N4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- F- k4 C" b$ B; {4 D
, M0 p1 r" t' b7 R+ vnum_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
) w$ i$ F' a; K( E
# 特征提取器
for param in model_ft.parameters():
    param.requires_grad = False
1 w; l) `* ^( _; ^4.2 修改成FCN
) E9 W6 {$ N) H- i. U
0 Q$ m; ~- U- _8 c6 l$ T全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
. y9 x+ A0 _8 F+ }' A9 d' W: m最终输出的通道包含了该空间位置像素的类别预测。
5 N, j8 B" }4 j- S. S
: j$ S' m. r$ l# W8 X对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
& ?* U. e4 c  G. ~- d% p- G
9 @  T9 W9 {. @' r可以先打印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], # 去掉最后两层
              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) # 转置卷积层使图像变为输入图像的大小
0 N: E5 `$ g* @9 P
# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
; F5 X  d, b% L" V6 o+ f" Sprint(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
& ?. a$ {( h" b! F# for X, Y in train_iter:
8 f" L( n" g2 X3 T  H9 l( o: d) d) G2 L#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
#     break

9 E/ L3 Y( D( g7 g( @$ q/ Q  h- U
4.3 初始化转置卷积层
0 W% u  F& L! v/ ?4 s- u7 \
在图像处理中，我们有时需要将图像放⼤，即上采样（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函数构造的卷积核的转置卷积层来实现。
( I7 ?' H" |% m/ j) c" p4 j
, |" L2 Z. \, t8 E1 \* F, k
2 c4 |8 O& B# d6 _( y( N; u# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
  r2 _( B- l( _    factor = (kernel_size+1)//2
1 Y; C4 p  G, I6 L0 j! `' D    if kernel_size%2 == 1:
        center = factor-1
    else:
        center = factor-0.5
2 [9 B( I2 ]* T; v: G" [% u    og = np.ogrid[:kernel_size, :kernel_size]
0 L  y$ r8 h3 j    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')
4 A+ s8 R. r+ \# i) a2 b8 ]    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
$ K; t: U  e& t; }& _    return weight


% F7 ^1 z% D2 V- M" Q在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
/ c7 k+ s6 w& m+ w6 K+ p+ [& I. A' b* u7 t
2 X# t) q- D* ^9 ^+ R) Unn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
' }+ M" D  P4 t8 \1 I! S' X1 Vmodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
* r) G5 C% N3 }# v- s
1 D1 r# ^- O0 U3 t! B- b
$ Q! q$ q8 R: F( j/ O! l  ]
5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
( ^  U  Q& [# N2 `    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
. j/ x" p  F& X; ^0 w  g! @    best_acc = 0.0
# T8 o0 \# ^: |+ g# ^8 f    # 每个epoch都有一个训练和验证阶段
$ x" N+ M. U/ p- H) i: I    for epoch in range(num_epochs):
3 {6 \3 {, q0 d0 ~" {        print('Epoch {}/{}'.format(epoch, num_epochs-1))
/ q, B3 r& c6 L  C2 T) h+ G* h        print('-'*10)
/ e( C1 X$ I0 @- X  R3 b0 q        for phase in ['train', 'val']:
            if phase == 'train':
, d4 u) D6 [7 |                scheduler.step()
) w: l- H' e. [* c                model.train()
" E2 v, ?; R: @3 C1 n' j! q            else:
                model.eval()
            runing_loss = 0.0
            runing_corrects = 0.0
+ W! f) ]- t1 q. W( g1 a$ v            # 迭代一个epoch
( ^  q% F, }. q4 a( C2 C# ~/ h            for inputs, labels in dataloaders[phase]:
! p. ]! V9 W2 Z5 ?% ^' W* N                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
4 y, S& o% M) u) H                                # 前向，只在训练时跟踪参数
' U7 [1 c- x  u' l3 q6 a3 J                with torch.set_grad_enabled(phase=='train'):
% N4 R* c- j$ V$ w                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
" p4 H4 h& ?9 c& q  I) T0 @9 ^                    if phase=='train':
3 ?0 d( ?. V, v& }2 Z0 u                        loss.backward()
) B6 r: v6 C; z: j& l( a  F% ~( d' o                        optimizer.step()
                                # 统计loss和correct
4 }4 n9 D2 s; N- T                runing_loss += loss.item()*inputs.size(0)
% V* r& p1 x/ D- }# G3 @; k                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
$ b2 ]0 s) V) o" C$ p& X
            epoch_loss = runing_loss / dataset_sizes[phase]
/ Y( P, K8 g8 a# p2 ]* g9 z- ~            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
* U5 K5 u+ G8 X$ f1 `5 U7 f            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
- d7 U6 I1 [# E; ~' B! O5 n% j6 W                best_model_wts = copy.deepcopy(model.state_dict())
        print()
    time_elapsed = time.time() - since;
# U/ W" E) W& {    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
) ~$ U) U; h* r3 m; y    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model

# m6 j( p. L/ i$ ^下面定义train_model要用到的参数，开始训练
+ ?# t- `; e0 N+ ?  m* N4 u: y# v
( L5 E( Z( M' l; ~) s" ?epochs = 5 # 训练5个epoch
  t* z: ]/ k% ncriteon = nn.CrossEntropyLoss()
optimizer = 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)

$ A3 `+ w* {$ o6 S1 A# 开始训练
" N4 T! S+ T. c4 u) b& omodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

# [" a4 n# y1 l+ K5 j: T6 测试模型

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

def label2image(pred):
5 F+ `% M" E5 Q3 n( I8 M    # pred: [320,480]
% c, e; [2 X8 ~" s1 [, ^2 k) S    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
7 j' B  V3 {+ c+ X8 f) v    return (colormap[x,:]).data.cpu().numpy()
. o" H: b3 o# B1 M2 i& q' {
9 [# D5 D9 x$ T; T& S- L

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

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):
' J1 X+ J$ v# R$ w    was_training = model.training
/ Q1 H6 O+ M' D9 R! g8 j    model.eval()
2 _1 u4 M. I# t" c6 H2 @. b    images_so_far = 0
( \: T" b2 Z; h1 q    n, imgs = num_images, []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
- {2 \! @/ c% ]: n2 T) h            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
2 ~& t" ?2 F6 j4 Q) a$ ^* y! K            outputs = model(inputs)
            pred = torch.argmax(outputs, dim=1) # [b,320,480]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦

! |+ E9 ]: J7 S7 i  e8 @            for j in range(num_images):
9 z3 C& \% R* p- B2 _  l                images_so_far += 1
                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
5 _* k' w0 `2 V% O" S                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
                    model.train(mode=was_training)
" e7 _; L9 ]7 i' p: L6 L/ Z' s                    # 我已经固定了每次只显示4张图了，大家可以自己修改
2 I' g+ |: [% l                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)

# ~  u; o% r& t' F- X# 开始验证
% N. S# ^, A9 `* ]! rvisualize_model(model_ft)
) W; m2 v, n) M9 E& v& X! I1 u6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
! n, H% T2 {) \9 |def predict(img, model):
    tsf = transforms.Compose([
& h) P" }4 N2 r9 p( F0 Y; y  f            transforms.ToTensor(), # 好像会自动转换channel
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
! v* V" y0 s% t" h& X( g    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
  m  n% ^: V/ \    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
. \- w+ _0 x: C- _2 ?" [' J, j
def evaluate(model:nn.Module):
# m- R) K9 j' [7 q7 E7 w    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
7 y/ f3 O+ z! r* `( m    n, imgs = 4, []
    for i in range(n):
7 L- P/ e. b; C: f8 `+ g% l        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
3 D$ {% ?# D2 r3 O  C0 p. a        pred = label2image(predict(xi, model))
. y* ?$ d# a! s5 I. y  i        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

6 w; G' x) r4 o6 [# g( x- z# 开始测试
& B, D  N* E5 @, q* t$ xevaluate(model_ft)
/ i5 h9 {4 Z, y, K& f
4 q0 P; ?  c: i$ |7 p7 结语

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

+ e/ |3 {4 w+ v2 f7 ^  nEpoch 0/2
----------
1 G6 O) w/ g/ U, I7 ^* ttrain Loss: 1.7844 Acc: 0.5835
1 q+ \4 i- @2 ^+ y) ^8 Ival Loss: 1.1669 Acc: 0.6456
2 ]- M+ G. [5 l5 T, p3 J8 F& a
0 }) P  ~5 C- h8 a! sEpoch 1/2
8 C; o& ]5 X2 n( s6 ^1 r# s# F----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929
" ?: Y! v' W! q2 p! J6 c
1 ^+ ^3 F' c9 m  ~Epoch 2/2
; o  R) S" k2 i  d2 K----------
6 t& g; u/ s5 c; p- c/ Jtrain Loss: 0.9578 Acc: 0.6706
4 |' H, a$ X) Z' l1 \: c9 j& Yval Loss: 0.8088 Acc: 0.6948
3 x' m, a: A* _# P- u) K
" |! w, P" [7 O7 X$ f' eTraining complete in 6m 37s
/ P* E8 k6 O5 Q1 ^0 L' m



当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

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

$ E+ C( ?4 e; o: Q+ z3 g  x
语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
$ o' F' h/ B1 ~, j
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
( X; k! @( t$ K1 ]9 t% Z" M+ K, |GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
* S/ l4 r! `% f, ^ENet 用于实时语义分割的深度神经网络体系结构[论文]
" J+ b) w, m: OU-Net 用于生物医学图像分割的卷积网络
" E9 R' k1 ?2 n5 Q7 S) wSegNet 用于图像分段的深度卷积编码器-解码器架构。
4 s4 v: m5 `* K+ \还有（DUC，HDC）、PSPNet等。
0 H' J: c$ A' B
常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
7 K6 ~5 h+ K6 l% O) T9 K, J
5 P4 o1 S/ l# w/ ?6 ?9 ?对于损失函数，除了交叉熵误差，也可以用这些：

- E4 _8 l- B3 H3 |& H" y) FDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
' k5 \- i% A' ?: p( N" C6 nCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
" h. A  k& p9 ?Lovasz Softmax 查看论文：Lovasz - softmax损失。

5 L6 s. ]' h  T

# \8 [8 Z6 l" R8 i
7 R5 @  O! o) h1 b
! C1 a; j: e: C; B6 F% K8 i————————————————
5 F' h5 k9 E# Z  o3 \/ Z  z3 b  T4 T' \版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507

5 f$ `, j, h% _9 q( u! h- ]& @
% S1 P5 i1 [7 i4 v