Pytorch实战语义分割（VOC2012）

杨利霞

; b# m7 t& M3 f% Q4 d1 W
" w% l7 ~: w( y; d7 Y  [Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
, i+ V4 A) v) ~; B3 x
3 y+ ~$ B/ w# M. X
3 r( @3 S& o/ K: }语义分割中图像有关狗、猫和背景的标签
  U" _# z" S& n* U& ]  H3 Z文章目录
6 h- J$ C% X: C) F
1 图像分割和实例分割
+ Y1 @% @& A+ t5 A; K2 Pascal VOC2012语义分割数据集
2.1 导入模块
3 ~' c6 `! M; R; s' k2.2 下载数据集
1 }) }; r! J8 `) B; G4 d2.3 可视化数据
" w, N1 N& g1 T" M2.4 预处理数据
3 自定义数据集类
: ]6 U+ V% A; U/ L) c3.1 数据集类
& |% G7 c% {' J5 R  w* L3.2 读取数据集
# M* Q6 c. ^% c. b8 E) H# k: @4 构造模型
4.1 预训练模型
4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
, O6 J, U' {8 o4 d, u6 测试模型
8 {2 l$ f% ?/ ^( W% s3 Q7 O  N6.1 通用型
6.2 不通用
7 c7 a7 S0 U5 [8 D, e5 c) C7 结语
2 C& C% ]% d% \% h& u% }1 图像分割和实例分割

4 j% u' n4 M  B5 h  U/ [7 m# l计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

2 p' F& @. T; y+ g0 h图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

# U0 g4 Q1 K1 }, G$ b2 Pascal VOC2012语义分割数据集

. s) C! ]* J" v5 }7 W2.1 导入模块
import time
import copy
import torch
from torch import optim, nn
$ S3 S& w6 N, _# H5 w8 ^8 oimport torch.nn.functional as F
9 s; x8 m/ A' Y% }- pimport torchvision
from torchvision import transforms
! R6 e' X/ Y! d5 U$ kfrom torchvision.models import resnet18
* `8 r3 d% p( c5 Pimport numpy as np
from matplotlib import pyplot as plt
from PIL import Image
import sys
4 y) x( W3 ?  b* C% vsys.path.append("..")
" b6 X, G8 X$ j: _- x, v- Hfrom IPython import display
from tqdm import tqdm
import warnings
/ y% Q/ @, p& uwarnings.filterwarnings("ignore")
- ?' m( C4 l, s1 x
1 |4 P1 q8 _7 M( P. k2.2 下载数据集
4 x- y4 p: I; i- [2 p+ d( n" z) d
: C6 h$ o1 I. {6 Y. F语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
. P$ i& J2 H4 x4 O8 m( Z# u

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

定义read_voc_images函数将输入图像和标签读进内存。
- V, `" i/ ^% f: g: B$ {
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')
    with open(txt_fname, 'r') as f:
        images = f.read().split() # 拆分成一个个名字组成list
. j6 L* i( Y' O& k    if max_num is not None:
        images = images[:min(max_num, len(images))]
/ E. g, P  k' {; p- X; P  X    features, labels = [None] * len(images), [None] * len(images)
. o( @' Z7 \% a4 `* E1 G3 Q    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
: K/ [- @4 y' U0 w& l7 ~        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
0 n- S' [3 Y# }. ]        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255
5 P, k. B( D0 R4 v" \7 K/ G
2 |; \  ^- l( N! W; G$ }: r7 D定义可视化数据集的函数show_images
4 c$ ?8 U) e. }
! h: j% j5 F6 p$ U2 n; o# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
( [7 c0 s: y4 N& J6 F    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
- O- x6 Q( \4 J8 L( I; M) T. Z) X    plt.rcParams['figure.figsize'] = figsize

) J! k( {1 {+ V( D! u, {- mdef show_images(imgs, num_rows, num_cols, scale=2):
9 j# M1 Y1 f, i8 f4 g; p    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
7 z( U2 }: k! C% r2 F    for i in range(num_rows):
7 ^7 Q1 n( Q+ W  m2 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)
' d# S* F" l  U6 u! _+ S  H    plt.show()
    return axes
  e) x8 _" P6 P! n/ |8 Q; ^1 L% g
定义可视化数据集的函数show_images。

" b3 q$ ^: a" e% X& L# 这个函数可以不需要
9 O( r9 O! B3 U6 D& ^6 odef set_figsize(figsize=(3.5, 2.5)):
% _: g% g4 S. P& r$ `    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
* \( z6 O; L6 c. i" v: C    plt.rcParams['figure.figsize'] = figsize
6 {/ ]/ V4 p$ t" E, |( Q! o% Y
% B+ l7 v, n( g: r1 C9 B+ [" Sdef show_images(imgs, num_rows, num_cols, scale=2):
+ }' c  \; `0 Z* r4 o$ c5 @: B    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
6 z. v1 I, v5 |    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)
* G( N5 ^6 g; z/ _            axes[j].axes.get_yaxis().set_visible(False)
" `2 A2 N# V* t0 V8 G: M  d    plt.show()
' Z$ A  j2 `/ f3 u    return axes
9 Y, q, ~' G0 [8 c* I+ p画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# 根据自己存放数据集的路径修改voc_dir
: o8 o7 k1 t$ J0 c2 z0 u" Jvoc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
. m8 H3 T. ]9 O: wn = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)



, E' h  X! \' J- w2 q列出标签中每个RGB颜色的值及其标注的类别。
% G) R1 n5 ]' a& H2 h5 u- E# 标签中每个RGB颜色的值
# R! e3 }6 m5 V& {6 LVOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
) I5 _7 F8 ~8 X1 i, ^6 I                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [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& }( o$ t4 R# w) X                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
7 i9 W" {7 D/ }9 t3 x' f& Y9 ?  Z0 l# 标签其标注的类别
6 [* S6 d" a- v0 S: O0 ?VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
6 m6 Q& X) }+ ]) W               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
1 I" }0 ~8 r4 n# ^) }$ C有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

2 R7 l& u; X. v) ?有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
for i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
0 @/ [  O# y! G  ]3 r& b    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
: t. z7 y' |3 p% C  }- ?    colormap = np.array(colormap.convert("RGB")).astype('int32')
0 ?: R/ j# S" \" S; C% r  \    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

可以打印一下结果

, [/ A6 @$ E: m! My = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
* s" F! Z% [8 Z/ E
1 l* X; c, v! A  Z2.4 预处理数据
1 ~: O# D' J$ F" c
; P1 k  t$ D$ x2 |2 l% F5 _在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。

& c4 }( R9 s; L4 a5 D; p
def voc_rand_crop(feature, label, height, width):
    """
    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
! f5 @! V* a5 T4 T8 U! V        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
3 h; j4 L3 y% N& @5 l    Returns:
+ l. w1 x; F1 `5 x        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
- u$ F, b+ e& Q    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

1 c7 E  _4 x  ~4 V& ]% |5 s% `# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
$ p5 E+ [- `$ L3 M! r; _3 x8 A1 }    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);
+ t4 z; T$ v5 A  Q# V: P  e) K


5 X2 J6 U; t" H7 M% {" |5 z3 q3 E" q

5 l' c% t( ?( J8 n
3 自定义数据集类
- o, X7 ]4 x, W/ g" N2 m; y' Z# N) G
3.1 数据集类
/ G" z3 a  T( d8 r$ a- X
9 h6 _. C# Y# @' G1 ?( _" ?5 E5 \torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
7 c+ b8 l- I% X  L8 N) U
( `$ z: k. w& Y__len__ 实现 len(dataset) 返还数据集的尺寸。
  E- ^# y/ V/ s% {__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

/ D' {7 o, j- s2 |6 J! uclass VOCSegDataset(torch.utils.data.Dataset):
6 V% }5 R- h% x! h! d    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
6 S( @: D' y3 O& d5 `7 w! U        crop_size: (h, w)
        """
5 S; y5 c) b$ Q1 c' E. ~        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
2 ?6 ?7 C# F' i" }. p9 ]/ O        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 b5 b; u- J6 l8 r/ o5 N        self.crop_size = crop_size # (h, w)
! ^" l, I- t& S8 H  |        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
. t, ?" J6 |4 o: H        self.labels = self.filter(labels)     # PIL image
& Y) M' c3 _# x) D2 }; {        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

$ f! A( J8 q+ _1 B" X4 a" ?, R    def filter(self, imgs):
        return [img for img in imgs if (
  O) P! X' y9 F8 ]1 s6 x            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

    def __getitem__(self, idx):
' K6 X8 q5 T9 W        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
# \+ m& X$ \. _: g, {6 m$ T                                # float32 tensor           uint8 tensor (b,h,w)
7 S- N8 X) Z" [1 R  R        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

& C; z9 [: v2 V' [0 t( k8 n; [    def __len__(self):
& ^2 G& I! f: H" }: \% F3 i- C        return len(self.features)
9 C8 F5 [2 ~5 N1 c) G3.2 读取数据集

2 D  O- Y) s. L& ?3 \- S, j通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

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

  U! I( {6 a- y# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
' {* H4 e. ^" r3 r& p1 K6 T
2 l3 l  L1 L) v! y1 u+ `# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
" n0 ?- ^( s" V' x" X, fnum_workers = 0 if sys.platform.startswith('win32') else 4
  g, N5 J$ v. W1 f* |train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
7 l! r" D6 z$ F/ \' M                              drop_last=True, num_workers=num_workers)
. b! C3 E9 E/ K) t4 `6 C* r. ~0 }4 j7 Ytest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
% ~0 I; V. ]1 i. _) S                             num_workers=num_workers)
# [1 T- ?0 s+ t) }
# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
9 y/ m1 P0 s4 @# H$ E. W6 Zdataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}
; C5 f5 N  p3 L, s- x# `2 Z: ~9 t# `
1 b; [6 n; ?' s# b( {4 构造模型4.1 预训练模型

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

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

num_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

! c' d. z- Y- m- X* o$ [# 特征提取器
; j% ?. j8 O! E2 F& ?  u0 hfor param in model_ft.parameters():
    param.requires_grad = False
% b8 \( U: O" C% l+ p9 M; S4.2 修改成FCN

; z: l7 j1 w- q! e& k全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。
6 Y3 Y0 y; X- K  [& b
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
* F, c) P5 K# O/ A. 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 倍即可。

/ I( d( N: f6 v' _$ m2 w& Kmodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
( ^8 U2 r9 q) 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) # 转置卷积层使图像变为输入图像的大小
5 u% w- g8 m, X8 _8 `4 N
( R" B3 l9 {" o! i3 |# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
8 `) T3 L& d& F
3 C4 M/ f9 e- X# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
' m, ?; b( |% t" c# for X, Y in train_iter:
#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
( [, k; u$ k6 ]#     break


4.3 初始化转置卷积层

- W- t; @$ ~0 M4 z6 b, x6 X9 A在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
- Q, b# M7 X7 D6 R2 L在坐标 (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函数构造的卷积核的转置卷积层来实现。
% h3 X8 e7 r0 M) P
0 X+ v1 U+ N7 O3 i: N$ c( j
# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
) ^. F( h. w4 }' n! l3 z+ P    if kernel_size%2 == 1:
        center = factor-1
    else:
        center = factor-0.5
0 x; K0 ^6 T$ q. ?/ D& c    og = np.ogrid[:kernel_size, :kernel_size]
    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')
  C) G! R' k; O. t& P    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
" j' K2 R) }1 O) o    weight.requires_grad = True
( c& q8 L, k/ F- O. c5 z    return weight

8 p& I% e' K) J% ~
% y  _; a! b0 @+ Z& A& b! d在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
: h+ s5 m, r9 v6 O  S) {
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)

% U& _# l0 \2 {, Y* O

/ Z" n) p* f3 T5 训练模型

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

! F5 D& e2 v: T8 cdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
$ k7 K: ~: r$ y% y) c% k    since = time.time()
, f: k% N. I5 w* }7 A. M- T. n    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
8 ~! U; F1 L9 Z$ Q( [% W! {# R        print('Epoch {}/{}'.format(epoch, num_epochs-1))
, S, p3 e8 m" d        print('-'*10)
        for phase in ['train', 'val']:
$ o! Q' m& u+ Y7 J! d: x; ?3 A            if phase == 'train':
8 l# Z9 e. Q7 A2 `7 V                scheduler.step()
                model.train()
            else:
                model.eval()
4 W8 F% @3 w& S* w# S( z) y            runing_loss = 0.0
            runing_corrects = 0.0
' l4 ]/ N( n0 H0 d' f            # 迭代一个epoch
' h/ p( u( E6 M1 z4 C& X5 U            for inputs, labels in dataloaders[phase]:
* w4 Z" l# h9 T  k$ @9 h                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
  W# H7 S- H* s% G( i" t                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
0 P9 r" D. J8 D3 G2 Z; B1 s                    # 后向，只在训练阶段进行优化
                    if phase=='train':
% ]8 m4 ^# [( M8 b                        loss.backward()
4 `; N* b& [2 z3 O. b4 E                        optimizer.step()
                                # 统计loss和correct
& X) P  I1 {0 G* w! g) T                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
0 ^" J; ~3 |( ]
            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
4 Q$ H6 p6 @# q            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
# ]& N& d8 r- o3 Y+ U; j% m            if phase=='val' and epoch_acc>best_acc:
: ]6 H- B# o/ I2 a                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
2 g4 Q' e9 F- a2 N5 S5 [        print()
" A$ l) l0 d# G& |8 U0 z( j    time_elapsed = time.time() - since;
+ y" A  q  i7 }& S2 [/ ]# I2 a% N    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
+ V% g* L& r. }. e    # 加载最佳模型权重
3 |% S' Q- t9 A1 [: r; x% x9 m    model.load_state_dict(best_model_wts)
    return model
; b5 R. G  f* B. C! ]( n
, t' {0 i4 e( Z" M* Y3 x下面定义train_model要用到的参数，开始训练

epochs = 5 # 训练5个epoch
. C' _, @+ C: Hcriteon = nn.CrossEntropyLoss()
% j* L* h/ x  k3 G( v5 Foptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
- M) q4 t8 d* ?exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

6 g# t7 Y2 P% A) c5 _% `8 s0 d6 测试模型

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

def label2image(pred):
8 d0 m! T! R' `% d    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
    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)
def visualize_model(model:nn.Module, num_images=4):
5 }& ]  {6 k0 d5 C' A, B/ y: E6 K    was_training = model.training
    model.eval()
, }, _5 N7 k: k5 r5 Y$ l7 s9 J    images_so_far = 0
% p$ i2 b( Z3 k1 Q: |    n, imgs = num_images, []
! V' ^; h! I0 g# i- z7 Z8 M. W    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
) X6 w4 S  E/ m( O. m* Q" i            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 # 记得要变回去哦

* l, |1 P; R: Z- _6 @4 g: F3 i" [+ w8 D            for j in range(num_images):
5 Q% l  l8 m+ D) ~9 R1 x                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:
                    model.train(mode=was_training)
# `1 D5 R# `( Z) D                    # 我已经固定了每次只显示4张图了，大家可以自己修改
" s+ m' ~3 `9 ~0 H! E5 h' p                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
; Y4 \, A6 q1 K8 f+ W( B+ W                    return model.train(mode=was_training)

1 H4 l. _3 ~, l$ K& \0 l# 开始验证
1 n" F2 B* i$ h1 d! Hvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
/ q5 h9 j7 _9 z8 |    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
% O8 l4 H' k, a9 G, w; b, H    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)

* u/ A: w1 ]1 o% S$ |2 O* Fdef evaluate(model:nn.Module):
    model.eval()
3 x9 E. }2 w9 d& T( C9 L    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
    n, imgs = 4, []
5 n, U$ m8 y+ Q4 o4 k0 S    for i in range(n):
        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
# p  T3 `) ]/ h7 t4 H! Q; f6 M% d        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

! B2 }7 N3 x  Q7 |# 开始测试
evaluate(model_ft)
' l+ J( d# a8 U* t0 R
7 结语

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

Epoch 0/2
' r# w7 n, t  N5 Y' j----------
train Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456

Epoch 1/2
6 V/ a- }, a  e8 C( n5 C0 {8 }----------
train Loss: 1.1288 Acc: 0.6535
; w0 S5 u5 D6 Qval Loss: 0.9012 Acc: 0.6929
& u" w0 K( F( V4 n/ a  l7 Z( D8 b* H
& c; h  d4 [& k* hEpoch 2/2
, N/ }$ i0 l, U9 z. v+ R# R----------
7 H) N6 g) @9 j% V/ f. xtrain Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948

Training complete in 6m 37s
/ f4 B+ x/ p, `' m$ Y2 m
, n7 }' o2 E/ {% x" j2 G! l) |, a
5 {2 b1 ]# `6 }4 v8 a! i

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

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


语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

7 u# _/ U7 y. S8 A& V5 e% j9 Y6 ~( iDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
8 J! ?$ B: Q5 I* Z* g7 N! f  LUperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
; D( Z" a# s7 g  U. AU-Net 用于生物医学图像分割的卷积网络
8 a" R7 c, I; m; {SegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

  P/ I7 W: V% c7 i2 j0 O7 O常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
- h3 B. W6 f8 u8 x$ e
8 W1 h* B" N8 W% {7 e对于损失函数，除了交叉熵误差，也可以用这些：
7 w6 L* m) d% H( h( c) ^/ u
Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
' M. E2 s' J  S( O! B


" J5 b+ \' K4 a! e! W: [
6 W$ ^' K9 O7 {. X. f' z; f+ K! |
8 {& B0 @9 Z2 z, l2 T6 y/ F* `————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507

+ l4 L% |) p0 @; }
5 s0 V, m  H- L7 x$ N% T& Z: x5 h