Pytorch实战语义分割（VOC2012）

杨利霞

1 v- X: F! X; W3 J+ e5 [Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
) ~0 B9 |. Y: e% {; j2 y' u# \语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
5 ?3 [" |/ }; w* k
2 s9 ]% Y; L/ ?& t% g
语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
8 Y* O* q, o; j; {) ]/ q2 Pascal VOC2012语义分割数据集
2.1 导入模块
2.2 下载数据集
5 {* T- S0 y/ w( W, K( T& w2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
3.1 数据集类
3.2 读取数据集
# o: g- O) L. _; S4 |4 构造模型
3 E- g8 F+ n# ~* X9 W% h+ y4.1 预训练模型
+ T: A/ |9 p( t) W4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
6 测试模型
6.1 通用型
6.2 不通用
8 E" P+ W& |8 e" K7 结语
( d3 y' l3 v; e( x1 图像分割和实例分割

计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

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

$ w9 T$ |# G4 k+ p2 Pascal VOC2012语义分割数据集
2 c2 V, j; P$ z7 n% ?, d: W
2.1 导入模块
import time
import copy
4 W( d. Q4 G' g. Mimport torch
7 O4 q3 d( x! [: ]from torch import optim, nn
# o$ ^, d2 E6 t. S9 [import torch.nn.functional as F
. n3 r/ x5 U, t% h4 k1 B3 s. A$ h6 V9 {import torchvision
2 A% l4 f$ N6 s- R- Z* cfrom torchvision import transforms
8 B) G7 v9 @: U* Hfrom torchvision.models import resnet18
/ q( W3 j7 N( \( T8 @2 H/ Aimport numpy as np
from matplotlib import pyplot as plt
, U! u# X' a2 m1 E% v) `from PIL import Image
import sys
sys.path.append("..")
from IPython import display
& X* u: z6 }/ H( l6 q. D" zfrom tqdm import tqdm
. a3 l/ b- i% E/ W/ zimport warnings
/ M  S8 P5 Y+ m9 i# _1 n- ~/ R7 Dwarnings.filterwarnings("ignore")
1 z6 I. J. i; F1 @7 t3 g/ w8 L$ a
; v2 _* Z4 h7 f/ R6 E3 a2.2 下载数据集
4 t8 @, t. g: V, K
$ X8 |2 `# j9 q: a: j" h" b语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：


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

4 d  y: ?* v! r+ ~5 w  R/ Z定义read_voc_images函数将输入图像和标签读进内存。

def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
: b% }! H6 `+ S7 @8 z! E+ _  N; E    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
$ z$ E; f3 o( ]+ K/ X4 i    with open(txt_fname, 'r') as f:
' G, j/ |2 e$ ~% ]7 b        images = f.read().split() # 拆分成一个个名字组成list
  w3 i! M0 b" ?$ Y. Q# }: p; F    if max_num is not None:
3 ^  o# h# l: C( m, \        images = images[:min(max_num, len(images))]
8 k, u1 V# _' G- S& o/ ?" u    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
0 t" c8 m- N3 t; q/ r% @! Y        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
0 o. I& A4 x4 D: o' @& T. w    return features, labels # PIL image 0-255
) }% _6 v6 l. U5 [$ o
定义可视化数据集的函数show_images

# 这个函数可以不需要
* Z1 [$ ~, j0 Jdef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
% N0 G& m- N5 ^$ Z    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
& V" \* V+ b3 ]* _
. s* O9 ]1 g/ Z6 [8 _+ W  u* L+ idef 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)
; U1 @" g7 S3 ?5 n2 j, P+ V    for i in range(num_rows):
+ {0 w" H' @$ `        for j in range(num_cols):
8 R: g( C  x1 l  f; F            axes[j].imshow(imgs[i * num_cols + j])
4 w' b; b; p# N8 O5 ~8 K1 ?0 P            axes[j].axes.get_xaxis().set_visible(False)
3 _& k3 K, q! R& B8 c( i/ W            axes[j].axes.get_yaxis().set_visible(False)
$ ~9 P# Q; p6 m9 [3 H9 n' W    plt.show()
    return axes
8 Q4 l3 R+ c' M$ |
( Y* o5 L6 P: m% m; ]定义可视化数据集的函数show_images。
0 K- e8 N3 w" |3 u6 \
# 这个函数可以不需要
- b; ^3 N9 D8 @$ Y7 D* s% f& _& D8 wdef set_figsize(figsize=(3.5, 2.5)):
( A- g' A3 T3 c7 ~1 t$ B0 q    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
# B& p: n: n1 v    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
2 C1 i$ o! ]. m
( H" ^7 P' @  |" ydef show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
0 a5 C. e6 c' a8 w: N    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
7 b% s% c+ l7 F8 i, u            axes[j].axes.get_xaxis().set_visible(False)
! V( p2 A, N) [- e            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
. p/ J, }, h0 Q# t) T# L
5 I! o- _% [8 G4 i: P: O# 根据自己存放数据集的路径修改voc_dir
: ^1 c2 _* d0 D- a( Kvoc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
- G9 f5 l9 M) b, z4 g2 Jn = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
" d5 G  L- i& W2 @' b! W6 O6 Ushow_images(imgs, 2, n)

) \4 M8 @4 b# ^7 z
1 e- k. {6 e2 H7 y7 `
列出标签中每个RGB颜色的值及其标注的类别。
# Q8 |+ w7 d! w" d# 标签中每个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],
& z3 a& O% q- d                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
7 o2 D2 ?/ P& _, h# H                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
+ @/ ^9 \+ z: T: S                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
  @5 B9 }, a4 m0 {+ p4 k9 a. z+ m               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
# E% X" z) R4 ?% F' R               'diningtable', 'dog', 'horse', 'motorbike', 'person',
( G. V. i# l6 q# u               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
" E$ d( d4 j" @有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

4 v3 P& K: b, f: T+ ~有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，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
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
2 B& X' i( o) Z1 {* W# M% m4 o# e
$ `. c1 s, O: ], i7 t# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
2 B. J; k3 _3 z; ]    colormap = np.array(colormap.convert("RGB")).astype('int32')
0 _* w4 Z; C+ i  ~, q    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

6 G5 n7 T; Q+ ^9 O可以打印一下结果
$ l* ~' q) A3 e6 L
) a8 m! u" f' d. h  _8 A% Dy = voc_label_indices(train_labels[0], colormap2label)
9 k8 z" U; z$ j$ w; mprint(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
% K& ?/ B9 [$ v/ F
2.4 预处理数据

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

' _/ i3 A3 i$ q7 x; P+ o3 j: U6 _def voc_rand_crop(feature, label, height, width):
, ]5 R/ ~* c8 J' i9 L7 R$ ?    """
4 Q$ ?. v, h, i% q    随机裁剪feature(PIL image) 和 label(PIL image).
, G+ \! Z! w0 L& Q* i$ C& A, u/ Y    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
. i: k/ s: r5 C5 R8 Q+ y, r        img (PIL Image): Image to be cropped.
, j9 t* V! c+ G/ I. K$ c* r        output_size (tuple): Expected output size of the crop.
# F# S5 u5 |/ M    Returns:
+ z) s' }& I2 A. ?        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)
    return feature, label

% Y# |  r  ?; Z' Q# 显示n张随机裁剪的图像和标签，前面的n是5
; p0 W* D% a: q/ o( x, b4 jimgs = []
0 \# R/ `6 _4 f$ g  _for _ in range(n):
# N/ [( V; M5 c  a% l+ L    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
/ c* C( S3 ~, _show_images(imgs[::2] + imgs[1::2], 2, n);

9 `" j  e% _2 g/ Q
& V* L3 A! H5 D/ ~
* y2 s6 ^6 J7 Y- T9 M, e. w' Q3 B
1 T) z8 `* W1 N5 F% p5 ~8 t. V
5 M9 `' h- T+ {  L5 Q* D
# w( \. u8 }3 o8 Z3 自定义数据集类
- A( ~- w& a; n4 b8 A
3 _2 ?: L4 i0 y) K: a" F0 E3.1 数据集类

torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

__len__ 实现 len(dataset) 返还数据集的尺寸。
) i% H+ A) L. c, R__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
, J. j0 _5 y0 s9 ~
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
; i9 U1 f6 R( S, r        """
        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
/ O$ @' h* q0 Y% b" \! o% v        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)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
6 D: ^9 q3 ~  d4 c
    def filter(self, imgs):
$ ]9 f1 @3 k. C0 P        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

3 k' X+ E$ O0 M0 u; V  {    def __getitem__(self, idx):
- \6 b& q0 F! J& u        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
0 u& u. \* c! g! m! ~7 C                                # float32 tensor           uint8 tensor (b,h,w)
. I1 D8 U- f8 s. n; d        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

, G2 O, W$ q. u. F; G' f    def __len__(self):
        return len(self.features)
& w% {% n% i3 ]% }% c3.2 读取数据集
; b/ S; Y7 q1 i' o! d: ^! v
8 y* M" ?; V0 T0 \' n; D通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
8 l% Q+ `+ X2 a) b* [9 t
5 V9 t; k, u5 `' o# @, Wbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
+ X! i& F! U! B$ Pmax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
, A3 F" Z2 N7 R7 e
# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
9 W& h& t0 m9 u- m: D  `* s
6 W) p4 q# A0 y; L0 ~. e# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
& R; r/ V$ D, p. |% Knum_workers = 0 if sys.platform.startswith('win32') else 4
: D- A1 R' Z, m6 R4 Y" g" F. s2 h$ R2 dtrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
' ]9 n( q3 e# a( N) k4 rtest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
                             num_workers=num_workers)
& E: s$ w: j$ q% ?8 x  n" r
# P% g( }; |1 `; v$ h# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}
/ {% E- ]6 ^, Q% U! f6 M9 a9 |. p- D
4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9 w' e% V# K* a
# v6 n/ [& w& ]) D7 bnum_classes = 21 # 21分类，1个背景，20个物体
model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

+ q# L- l# a* [* }% {# 特征提取器
8 m9 ^! c$ ]' H* Y) u" Afor param in model_ft.parameters():
1 X0 [* z( O4 i8 ?) {8 B( ~+ r    param.requires_grad = False
4.2 修改成FCN
& O9 ~. v4 C1 w
! a: k. a/ Y8 Y" I) w全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
* P+ }8 [. d, C# V, S8 }最终输出的通道包含了该空间位置像素的类别预测。
/ _8 H) F# }" h
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

& v7 t& s7 w$ U7 |9 C1 C* ]可以先打印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 倍即可。
( r  A! W8 ^# N# I7 l  b' \% o% u! t
  }% {1 M: S9 Q7 Bmodel_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) # 转置卷积层使图像变为输入图像的大小
7 m1 k0 a4 _2 i9 B7 U6 A
# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
9 C/ y2 C* `8 ^% F+ V/ S#     break
# ~& s( U  Y( }. d% `# l
! U% g, i, Y( Y
4.3 初始化转置卷积层
( V8 b0 j- H+ D. _  A% X
1 _- u# ?$ L$ G. R3 P8 A4 Q在图像处理中，我们有时需要将图像放⼤，即上采样（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函数构造的卷积核的转置卷积层来实现。

9 b; a% Y! x; p0 x. G, ^
# 双线性插值的上采样，用来初始化转置卷积层的卷积核
1 e- M7 }& p7 wdef bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
$ K8 J  c% c6 G( E( k        center = factor-1
# Z" q  }: x) E( P. U5 Y    else:
6 F! V" ?; @7 V9 A        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
+ q3 ]9 ^% @6 }4 z+ e, w5 t' J    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
2 w. ~8 _3 J! }5 U. w  z    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
7 A! W! c# \8 H! `4 v    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight


! w# B4 J- K1 B" z% }6 L6 o; b1 W在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
6 I+ s- h9 o7 w% c, [+ H- U  e3 A; Cmodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)



. m+ a8 n0 d& W( c8 ?5 训练模型

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

# E6 f4 b2 q& i. _def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
! \. F  Q9 h4 v, [    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
! s* B1 m( y5 M8 M    for epoch in range(num_epochs):
& U* h: N% ?5 I) [        print('Epoch {}/{}'.format(epoch, num_epochs-1))
; X. v% h  T8 c+ [1 \        print('-'*10)
        for phase in ['train', 'val']:
3 v+ o, `; U9 a, }" E+ q            if phase == 'train':
                scheduler.step()
" H( S0 A0 [, {: L. W# G4 ^# z                model.train()
            else:
3 j. V1 F) y5 |                model.eval()
2 ?1 e7 ^4 Q- L            runing_loss = 0.0
# z' C$ m- \4 C+ p            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
' [( `. H0 f- \: y5 ?+ E, `: B                inputs, labels = inputs.to(device), labels.to(device)
! F" F( O* ~$ {, p: H1 ^                optimizer.zero_grad() # 零参数梯度
8 P" E1 }. n+ g* s5 B                                # 前向，只在训练时跟踪参数
% k5 f+ G, D1 ~: L* y                with torch.set_grad_enabled(phase=='train'):
2 n* ?" n1 D& w& l" X3 q- l" f                    logits = model(inputs)  # [5, 21, 320, 480]
, g1 E; T- J7 G8 `# c$ Q9 ]                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
/ i4 g7 x, L& Z# M                    if phase=='train':
. L$ P, ^# c9 j0 o                        loss.backward()
& m5 ?2 c. \; D, k9 V- b                        optimizer.step()
) B. H. Y0 k, Q4 Q, j) k                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
3 U( X* F% S; }* @4 U: ]6 S                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

$ Z) `! L! u' p            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
( G( k2 `5 w" F9 X- B0 [                        # 深度复制model参数
" m+ o/ l* e. e* `& ?. J  t            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
0 s' b1 q: h9 ]; d" Q; h, U                best_model_wts = copy.deepcopy(model.state_dict())
        print()
  ?- n/ W* V4 z    time_elapsed = time.time() - since;
& N3 Q8 m  a" b% s4 y0 F; `* B    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
, e1 D7 P1 M' M) I8 i    return model
; [/ I  O) o) |+ b
下面定义train_model要用到的参数，开始训练

3 q0 J7 A" o  n! `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
# {$ N/ |7 P- j: W# [# Texp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

6 C9 q( v( m" }9 B, I( S) \! D0 W# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
* g2 M& O/ _7 s4 c# A) D% A3 _
- M( ^- k1 N; Z  _! L) X6 测试模型

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

def label2image(pred):
% K, D* {: a4 ^: b- D    # pred: [320,480]
6 i4 A9 C3 i; a& s    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()
7 ?$ a' v$ Y& z3 A) V3 `

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

6.1 通用型

其实如果要用于测试其它数据集，也是要改动一下的 : ) 😃

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
& R! r; a+ @- [3 A$ \& y) ^0 [std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
def visualize_model(model:nn.Module, num_images=4):
( z( M0 D* E5 l; h8 _4 e3 g    was_training = model.training
8 U! U" m) E4 K$ R) g8 G    model.eval()
. ~" v2 T7 W* Q6 n3 G# A. ?    images_so_far = 0
7 g! q; s/ V& z0 I9 v' h/ ^" F+ @    n, imgs = num_images, []
( i8 ]- G+ U: _/ Y/ W    with torch.no_grad():
  R  v" N8 d8 ?6 B0 A6 E  d        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
) V0 [  y/ W0 J            outputs = model(inputs)
$ D2 n& r6 G) P. N# T9 G            pred = torch.argmax(outputs, dim=1) # [b,320,480]
* p! ]2 A0 e  m8 a            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
1 p0 K# h7 h) V$ I; v* L4 S
$ H2 `9 i8 G7 [3 ^. K2 S$ s            for j in range(num_images):
% _- w* U* @. b2 }, a                images_so_far += 1
0 W* _2 H, m5 F) H; [- G                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:
1 a/ ?+ b2 P- l4 f8 o                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
. T4 y5 ]+ a5 n: Z: a                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
0 R$ Z# m+ k7 C/ w& k0 M                    return model.train(mode=was_training)
9 ~# a! D. I/ b
: P! S) v& }% e$ t# L; K, S, x5 d# 开始验证
0 q9 N6 J. a9 N+ @" F/ ivisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
1 }+ W9 m  @8 a' L- A" g6 x8 edef predict(img, model):
    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
" G" F5 \6 N0 [" L! \: C            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
( @; D3 q8 r  b$ a
def evaluate(model:nn.Module):
    model.eval()
: F8 ^) k3 {6 V7 f* t7 F1 ]    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
0 E# f( T, m8 x) P$ q8 Y    n, imgs = 4, []
) D; a, a( b( }. ?9 m$ O; f( v$ F    for i in range(n):
        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
$ s5 @0 G5 g( w) c( n7 J2 v        pred = label2image(predict(xi, model))
$ U) u; ~" O- j! L. i& _( j        imgs += [xi, pred, yi]
3 l3 ^# e' m0 j8 N: c* h    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

& f3 `7 c* I. j* Y# 开始测试
( L  c8 ]4 @8 _& w& X+ levaluate(model_ft)
6 H' l3 l! e0 H! f
7 结语

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

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

% `, C0 Q' p% h, m) ~5 a- d2 m+ |Epoch 1/2
; }; j; b. l. O----------
train Loss: 1.1288 Acc: 0.6535
1 y7 m- s& Y4 f6 yval Loss: 0.9012 Acc: 0.6929

0 `! d, v8 v" o# p' k( \5 jEpoch 2/2
----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948

: b5 Q) T3 s" `0 {# tTraining complete in 6m 37s

  _1 {9 X/ U% f% j& o

( C! P- U( [( V
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

, B5 B3 j7 |- V0 h/ y- v8 a对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
# e8 `+ r  n* V3 g
9 S) c, M' R# U2 N' e& _
: k1 C# ?# F1 M  v8 U) w1 z! w语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

0 n) Q6 I' ?6 J, w' qDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
/ c, o) i7 w2 m$ X; _; ZU-Net 用于生物医学图像分割的卷积网络
5 e7 w4 a5 J  v& rSegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。
4 h& N: G; ~4 r
常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
  P& ?) |" x5 T+ ~8 d
, ^6 P, b: m1 `6 s+ R对于损失函数，除了交叉熵误差，也可以用这些：
9 h  |1 K% j$ e* G( n* c; q
- `, a' u. `  q( g4 [* w+ ]Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
9 {+ h" y& P1 _* V* ?( q9 yLovasz Softmax 查看论文：Lovasz - softmax损失。



, s  ^( K6 y$ q5 ?* ]

————————————————
8 @3 o5 m! Z! t; Z' X' x9 i6 M版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
, v6 l* T4 M2 y9 X
) @; _1 Q- h% E& c' k! @0 q
2 ~2 `1 G8 f/ z- u* z5 A