Pytorch实战语义分割（VOC2012）

杨利霞

. G* F9 m, Y4 P) h" [
Pytorch实战语义分割（VOC2012）
8 o+ v0 n2 o$ n( j$ X本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
2 Y% C, b# A. [8 [7 v! r8 j* v8 b

语义分割中图像有关狗、猫和背景的标签
文章目录
* W/ B9 {7 F: Y; g
4 G" o9 `5 a9 H1 {* |: ~/ |1 图像分割和实例分割
) J, q9 v# Z: ~2 Pascal VOC2012语义分割数据集
2.1 导入模块
& Y: P/ U/ g0 L1 _2.2 下载数据集
3 t- W0 Z9 ?/ S: c" a2.3 可视化数据
2.4 预处理数据
: Q0 [" I8 O- V$ v2 Y* v% h7 t3 自定义数据集类
3.1 数据集类
3.2 读取数据集
# Z+ E% v  H$ E/ U4 构造模型
4.1 预训练模型
- j" X. y' }3 S, k% V& N8 C4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
) Q( N/ |% S$ u% M/ q, d& e6 测试模型
) `1 b. k- F7 J$ ~6.1 通用型
6.2 不通用
7 结语
( N! p% ]& o+ x9 d7 c- C1 图像分割和实例分割

计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
6 U1 ?, T4 x1 E
, w5 F( [& \) {图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
& u  k( x2 e0 [; `* S) O+ X实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
2 G5 H3 H: T/ x/ a: c( _3 V
' X' w( J" a( M% k6 o2 Pascal VOC2012语义分割数据集

& c; }. V' m/ E) h- B- A2.1 导入模块
import time
6 X! W' k# y* d! w& P3 ~import copy
import torch
from torch import optim, nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
5 S, F1 d4 \" T  t( |6 t: ffrom torchvision.models import resnet18
import numpy as np
! f8 _. k) t/ s1 D0 Z( zfrom matplotlib import pyplot as plt
: u3 B* R; k) O! O8 xfrom PIL import Image
import sys
# Z4 o5 {3 v% t2 S9 S& `8 Ssys.path.append("..")
2 |! P  S* {7 {5 s; B. q* wfrom IPython import display
from tqdm import tqdm
1 r: N0 k6 F$ G2 qimport warnings
& u9 D) L( s/ D' F1 [warnings.filterwarnings("ignore")

5 ?  A5 C4 O% z& D9 ~/ B2.2 下载数据集
9 [) i; M4 r8 J$ J
( t- K9 i* E9 B+ z语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：


ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
- s0 m) Y) i6 n" J- j7 V& z8 Y4 Y: X2.3 可视化数据

' E& f5 w% [5 n定义read_voc_images函数将输入图像和标签读进内存。

def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
1 {9 G2 ^" ^2 ?    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
    with open(txt_fname, 'r') as f:
, d# Y. I. X7 g7 Y* S$ k/ ~, `        images = f.read().split() # 拆分成一个个名字组成list
) k5 y+ s& @. ~    if max_num is not None:
6 Z# |3 n  o  J& n0 _        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
( O7 X5 j0 s: P- ]    for i, fname in tqdm(enumerate(images)):
0 x/ [; r  k! Q- F7 Z        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
: o% u& h+ _2 }* f        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
/ l- ~: h& |8 e6 R. q6 O    return features, labels # PIL image 0-255
: m& B9 F" r  m
定义可视化数据集的函数show_images

# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
8 F6 ^) T2 T" ?* f    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
) o2 w9 b. G2 ^& s. X+ S) V# J# u4 f    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
' Q- o0 k. _  }( S6 z& S
def show_images(imgs, num_rows, num_cols, scale=2):
) |7 V! \5 H% l1 i- c    # a_img = np.asarray(imgs)
  ]/ g2 \' o8 z/ L7 L% C0 j7 E- ]. N    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
, o$ E4 \( G" k' L    for i in range(num_rows):
        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
) I  e% j- k8 E; d; \9 Y
9 w9 Z6 F$ h/ Z) s0 u. K定义可视化数据集的函数show_images。
# w! Z& ?% l) k
2 L+ f4 @. P- m, Z/ k7 q# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
1 Z) {. r- h0 X" r) E$ M- E    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
" F! Z: k( w6 y, }# q( h
& \3 ]6 p' @; udef 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 N) D4 _* ^" O' n    for i in range(num_rows):
        for j in range(num_cols):
& J, D6 M; v9 ?3 \            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
: S( a3 @% s+ N$ ~) X            axes[j].axes.get_yaxis().set_visible(False)
6 I, I9 V, F, [0 _- M    plt.show()
0 n& H' s- c# k( L, K; ?    return axes
- U3 R6 b8 P: e1 O) X3 I8 j画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# g9 h1 d  I! o$ n# 根据自己存放数据集的路径修改voc_dir
  C6 ^: e6 G; E0 V4 Z# f, xvoc_dir = r"[local]\VOCdevkit\VOC2012"
# U; h3 f# ~7 O3 u7 |5 m' G7 `# @7 dtrain_features, train_labels = read_voc_images(voc_dir, max_num=10)
& V' B0 [! i, ]- Y0 _8 _# M' ?n = 5 # 展示几张图像
4 p0 e% B6 @- J( `1 F' Gimgs = train_features[0:n] + train_labels[0:n] # PIL image
+ ]$ p9 b/ Z: Hshow_images(imgs, 2, n)

1 |, _' N& @/ f2 B/ }0 Y+ r

7 i& i" d2 O$ J' A" m列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
) z$ @7 ]( w' X+ x5 n. |  E                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
# |7 F0 I. I# u# Q5 J                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
! h: L' i/ ^* i                [0, 64, 128]]
# 标签其标注的类别
) N( M; g) b( j1 E% n" K( H8 MVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
6 ], L! ]" C! ?; W; A8 [9 S- O               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

+ G' _- I; e% ^) m8 X有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
! ^$ z; z. l& b, l' Bcolormap2label = 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
; c/ {! o  K/ w  M" Y, [: V. Q
# 构造标签矩阵
def 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中计算的下标

+ q% d. @6 Y/ ^) T可以打印一下结果
6 H5 B7 f9 G4 Z; r0 b, u. v# g1 W
y = voc_label_indices(train_labels[0], colormap2label)
( g: u) [/ T% c+ r; Q4 |print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据

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

7 D* u) L0 B6 M* j) e7 S* t7 Y% m/ F
def voc_rand_crop(feature, label, height, width):
    """
) w7 g+ Q; |/ X4 }  ~% ?    随机裁剪feature(PIL image) 和 label(PIL image).
9 C1 u. B( N. Y0 a3 {' s    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
4 o' b: `; d- _5 q    Args:
$ U7 `  w: k; C3 A, j9 T7 g7 _( {        img (PIL Image): Image to be cropped.
- v% D$ A( Z1 G2 Q" o        output_size (tuple): Expected output size of the crop.
' Z6 C0 b& D5 n7 N. g, w- o# Q( g. {$ L    Returns:
: H' W- h6 I; I/ @0 T        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)
& G  @& ?6 M4 m. ]0 {    return feature, label
: d: G) F( E) {9 _6 p
, r9 v) p8 N: H7 h1 X0 [% h# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
' I% I, a/ I% T* G4 i, U3 S$ w9 R8 Ifor _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);

" h( m; H9 H2 ]. Z! k


* _- ~. o7 I/ h+ Y
4 i( v% E! r' A( \
" X0 @/ O. C9 z8 o/ K3 自定义数据集类

$ M' |# z% j. d3 v0 c3.1 数据集类
' S8 w7 L% c$ J; ^' E+ y  C) _  E% z
$ t: _$ d4 ?! y9 ztorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

) k; j2 g- a9 p__len__ 实现 len(dataset) 返还数据集的尺寸。
# O3 P" N- g) G4 G2 x__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

( l- D3 E2 L; m2 l% n& J4 h  i9 cclass VOCSegDataset(torch.utils.data.Dataset):
8 |8 h- z& t# j2 p    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
$ N7 `6 v! O( o- I* M# q' y+ R        """
        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
5 h, `' h( U. s2 G; u$ r/ d        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
" U; V" [; W: x6 L( J$ P1 q        self.tsf = torchvision.transforms.Compose([
+ `' s6 q: ]! H) q$ E            torchvision.transforms.ToTensor(),
  }& H  W! t" i9 E3 Q            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
; V& c2 G7 c/ K) U5 _% R) l9 m        self.crop_size = crop_size # (h, w)
! u" Z9 J4 ]5 ^% \% W  z        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
2 t, A/ u9 e; ?1 e# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
6 t* |4 b. V" i2 z        self.labels = self.filter(labels)     # PIL image
! r7 ?6 ~/ v8 s/ P1 t. _% R) ?        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

    def filter(self, imgs):
8 N  l3 l0 F2 F# D) ?+ o        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

4 d3 X0 V2 \6 t2 h: r5 P    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
& l8 b2 J' }7 [                                # float32 tensor           uint8 tensor (b,h,w)
# I* j$ l/ H4 w- m5 a, L9 O        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

& e5 q, i. B% [    def __len__(self):
; {8 s+ _' U' j$ w4 g2 C        return len(self.features)
: F) O' ]: H) J* n) m) x" t8 B3.2 读取数据集
2 s/ d" o- u, n$ |& }
通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

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

# 创建训练集和测试集的实例
& T1 X; I, ]- k( j! J  R  l1 }+ H3 svoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

* Q# P# A- E1 l, k# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
8 x: y& n) b& gnum_workers = 0 if sys.platform.startswith('win32') else 4
train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
( O7 Y) g/ Z" ~) H  M/ K                             num_workers=num_workers)

3 l# |/ b4 k5 L' C" |# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
. ~8 [: U) d. y6 H& O* [& Hdataset_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')

num_classes = 21 # 21分类，1个背景，20个物体
( N6 e: _) O9 q" w! H% v8 V  v' ]4 R8 fmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
" H$ _) `' ^- R( ]/ U
$ g- X/ Z2 s6 `" [# Z# 特征提取器
for param in model_ft.parameters():
    param.requires_grad = False
4.2 修改成FCN

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

& d1 J7 `2 y2 D3 q3 W对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
8 I( O; e- Y" t- `! `- |
) X5 a9 E, Z$ r7 U% m8 F8 R% U可以先打印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 倍即可。
% ~9 Q7 R2 c: s, ?* u+ f
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
" }7 s* A2 m! f- m" c4 L2 ]" n; ?              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) # 转置卷积层使图像变为输入图像的大小

" `4 F- B2 Q; U) g% C& S# 对model_ft做一个测试
9 U2 {; e$ S6 px = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
9 f6 S6 p  s' ^6 u, g: |6 c4 ~. z$ p# for X, Y in train_iter:
6 _4 I7 n0 r5 q8 k( F#     print(X.dtype, X.shape)
5 t' [9 o/ E! F- t+ x9 p#     print(Y.dtype, Y.shape)
#     break


' J' j+ [' M0 d; P5 p7 S' P& G4.3 初始化转置卷积层
1 s. U$ h) I: E
在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
: K% D- d$ \& R7 S在坐标 (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函数构造的卷积核的转置卷积层来实现。


# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
! E9 W  Q& y2 e2 X$ Z2 c    factor = (kernel_size+1)//2
9 u8 `( d; N4 h    if kernel_size%2 == 1:
, B  }' O! b6 c% i0 M0 A  {        center = factor-1
: y& [, T7 X* ^( `1 L$ M    else:
        center = factor-0.5
6 i) l& k) R/ @# L! Q. Z0 z. g    og = np.ogrid[:kernel_size, :kernel_size]
7 v* B( h$ @* L& K! _    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
% M- f4 z+ E) T  q    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
9 ~/ z" w1 J; k+ A1 W) n    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
: R, S2 T6 _/ \! J& i6 m8 T    return weight
) e/ K/ P0 t0 T- B8 U" ?
- ^; t- l* r, ]/ ?  O7 E
在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
) m# G; g1 C0 [/ P3 |
/ s) ]9 u8 X. r& @  |9 gnn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
! l& d7 j2 d0 O, smodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
( i  E1 l( c- H: O
: X- H8 t# T0 o0 F5 l  e1 i) V

5 训练模型

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

3 f6 s8 m: `0 m8 \7 t: n2 ydef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
* `9 ~$ {+ n) g4 z2 X0 l6 j7 U    since = time.time()
- B; n" Z2 R4 S0 h6 _$ M    best_model_wts = copy.deepcopy(model.state_dict())
+ M0 v. Q# k2 |, u9 y    best_acc = 0.0
& N* S* }% e: p8 y0 ~& U/ Z    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
" D* B* [% {/ H0 B        print('Epoch {}/{}'.format(epoch, num_epochs-1))
8 {* `0 N$ S- t6 d+ e/ \9 B        print('-'*10)
2 P3 H; S' [# F9 d        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
( I# I- v3 W3 K+ h3 C9 F- ^                model.train()
            else:
                model.eval()
( ^1 o5 Z" U& \* ?' L            runing_loss = 0.0
5 y0 I. y' T( S9 o            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
5 Z" ^0 [6 _1 P2 r7 U# z                optimizer.zero_grad() # 零参数梯度
# ~% s; d( D, ~+ U7 Z+ G                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
* K1 r; y/ f! R# [% `                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
                    if phase=='train':
" e3 i+ [4 p0 u4 p1 x) H  l" n                        loss.backward()
& ~/ x/ t- Y# f* d2 m9 J; p3 L* c% [                        optimizer.step()
                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
+ d* R0 t7 f' V' r            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
/ U& m" W5 Q# |2 e& p                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
+ K: r$ g$ F' f  m                best_model_wts = copy.deepcopy(model.state_dict())
5 R5 M) k! m9 h  F% `2 d        print()
8 I* A1 _, c& G* A( J1 z/ _! Y# Q    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)
# j1 i+ e2 V. H$ B, u7 A1 S    return model

下面定义train_model要用到的参数，开始训练

) S! e/ T; V, A8 e' }epochs = 5 # 训练5个epoch
criteon = 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)
  _8 ~7 Y3 Z' l7 T' r& U6 C+ g* C
# 开始训练
; J* X6 O7 i+ F# Nmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

5 X- ^, \8 M+ s1 ~2 Q6 测试模型

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

def label2image(pred):
0 S- T  j" T# E! e, c, U) m    # pred: [320,480]
/ T% [6 _6 `7 d9 U% z5 q    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()

4 I0 z/ D" a  G4 Q

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

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):
1 q5 q. |! a: j; u4 ?8 S( q3 J- ]    was_training = model.training
    model.eval()
2 ^! E1 P9 }4 J5 R# {' m) r    images_so_far = 0
    n, imgs = num_images, []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
! o$ {. g# V0 O# ]5 M            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
            outputs = model(inputs)
  M- o6 B/ |: w6 P7 ^            pred = torch.argmax(outputs, dim=1) # [b,320,480]
8 f9 c( B% a% a0 q7 }            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
8 t  A+ E1 H4 |. u
  {% Z4 A2 D, z/ @" ~4 \$ c. B* b8 @            for j in range(num_images):
4 J. `6 r; g% ?" ^! W6 _6 N                images_so_far += 1
                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:
3 U( {" k* g% Z& Q                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)
% s* u. @4 W. W3 e
# 开始验证
visualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
5 A7 r2 x( o4 ^4 p            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
: R( E/ f6 R0 ^: A' n. P    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)

def evaluate(model:nn.Module):
    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
* K/ z0 w6 u$ O, b8 b: K. O0 h% d    n, imgs = 4, []
    for i in range(n):
# v5 R7 w* D( d3 m        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
! I* u6 _" N* [0 [% `    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
# |$ \% ]) x" m
9 T% y6 C3 a9 D6 O/ M+ _2 Z" ~6 g2 L# 开始测试
& W6 M. n. r9 S/ x* Sevaluate(model_ft)

7 结语

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

Epoch 0/2
3 Q1 g+ h+ a! R9 U! ~* K" ^----------
0 y7 a* G, ?" j1 i! b# `; o& e$ Jtrain Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456
2 P( U8 _6 z3 |4 i% u3 _2 d& _
( z6 w- Y( X1 V% ]- OEpoch 1/2
----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

Epoch 2/2
----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948

Training complete in 6m 37s


) Z2 p6 T9 U0 T( Z/ J% T

当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
/ Z  ^1 Z( ]+ W5 O: E' i6 [
2 F5 |. ]* Q9 N& b; K' U. d: z0 v对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
8 X( @! W' a) H3 w8 n
# [4 c- @$ z+ k- G$ \- u6 n1 h$ X6 x; g( x4 r
语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
8 [1 x0 P$ `6 y: z: [1 IGCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
' L' g' ^* {# Q; E7 WU-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
9 z1 L% d( G9 o. [还有（DUC，HDC）、PSPNet等。
" D5 l2 W1 ^- |9 [4 h; {( z
常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
4 B- x- A, Y2 H. b% D0 p% p( D
对于损失函数，除了交叉熵误差，也可以用这些：
8 g9 c* T# x7 |+ r. L. V9 I' x
Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
( T/ W; a! i0 K/ B* i' {* w" sCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
5 T& `8 v* A- M2 H- c" V3 rLovasz Softmax 查看论文：Lovasz - softmax损失。

4 y  y) c3 b, D4 a* Q+ }! `

, Y0 O2 V! [- j6 f  K

4 p8 y# {2 F, D: r1 o& [0 N————————————————
- ^! R/ f# h# ~( J, ~版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507

+ z( u# v& b0 z, h

: d0 n5 \2 V6 F- t. _