Pytorch实战语义分割（VOC2012）

杨利霞

. |& @# B3 }9 N! D* ^* T9 C, e
Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
" C! U, R1 u: m- n6 D语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。

, W8 [" E* h& s+ R
语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
" U! C9 g" l% ~4 \- ~+ K2 Pascal VOC2012语义分割数据集
2.1 导入模块
' Q8 U+ {( e; r; |  L" A2.2 下载数据集
% X; {# F) O; m5 V+ ]; i! _2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
: z  l4 D2 k# j0 M3.1 数据集类
$ k! e6 L1 c6 @, a3.2 读取数据集
! t: E$ w, _* f! x+ \4 构造模型
4.1 预训练模型
4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
6 测试模型
9 u1 J- F" o3 u6.1 通用型
6.2 不通用
7 结语
1 图像分割和实例分割

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

图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
4 |, c) y! v8 }; J7 V
0 n+ D9 s+ s+ p/ Z! ?$ Y& h3 i2 Pascal VOC2012语义分割数据集

2.1 导入模块
# i% r* ]6 \$ Jimport time
" ?+ z! f0 {3 ~" g0 ?8 Qimport copy
: w' f  a0 {$ yimport torch
from torch import optim, nn
import torch.nn.functional as F
7 ]. {' P! b2 j2 oimport torchvision
from torchvision import transforms
/ X4 M$ u! K- bfrom torchvision.models import resnet18
5 S: s, W7 y( Nimport numpy as np
! x& _  q6 I+ O2 Dfrom matplotlib import pyplot as plt
9 u! a& s& j4 _5 [8 Ufrom PIL import Image
: i1 b: ~3 f  x8 R5 ~0 oimport sys
sys.path.append("..")
from IPython import display
$ i% ]: y* N6 X9 c  n" Hfrom tqdm import tqdm
8 m+ w  L# {7 wimport warnings
warnings.filterwarnings("ignore")

* j$ W& K+ d# j$ k2.2 下载数据集
) F* j/ n" p4 K9 @5 E2 U
( W3 a2 R" v9 ^: ]3 |% `2 U% M语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：

; W6 k$ D  V6 {0 N+ p  r
$ E1 K% G, q' `ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
9 k  V5 U& G. C* D# ~( o8 T. V1 WJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
) s3 J% J$ `- e& a. y$ G2.3 可视化数据
8 p7 T3 w6 ?4 S5 A% h& L# q& p
定义read_voc_images函数将输入图像和标签读进内存。

) ^: ?' @7 z5 l- g% hdef read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
. X2 V, b0 h: I' ?) k5 F$ s  S    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
3 f$ s7 _, @# q, N0 j  ~' ]/ `    with open(txt_fname, 'r') as f:
& k- ^3 y7 b/ |$ ~        images = f.read().split() # 拆分成一个个名字组成list
    if max_num is not None:
1 B. i( V+ u$ m* _$ r! @5 A        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
3 v" U' j; s( M( ~        # 读入数据并且转为RGB的 PIL image
$ A: c" p' ?# A        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
" z" [5 `' s2 S" U; C  T% Z    return features, labels # PIL image 0-255

/ t6 p# u% k9 j定义可视化数据集的函数show_images
9 j( l2 V. t6 a
  Z, {! C$ F: ], G$ \# 这个函数可以不需要
9 X& S5 o' Y' @' t9 |. W$ c) Ydef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
0 M% M' ^3 m3 F6 V( p& w" t! M    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
# N* ^- Z7 s" s) z- i
$ S! Z0 [! C, E0 f6 ~, N; Z: gdef show_images(imgs, num_rows, num_cols, scale=2):
) o7 [8 w9 K# f' }* T    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
3 X+ J# F, ?2 k8 Y" m    for i in range(num_rows):
        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
3 O1 L3 z! O9 f$ c- k9 B0 ]            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
) K. p( j( D! @6 M' D& a! G    return axes
8 E7 W3 }7 }9 q6 d. h, G
定义可视化数据集的函数show_images。

# e/ b7 T1 ^1 k  o  Z* h; x# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
0 A6 `) d5 z1 m0 U. e- \    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
  Y; v* {5 r: y$ n
, y+ `. r, ]; a; tdef show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
! C9 A# ]1 z1 H- H4 D% L    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
! u% E1 s: P3 ?  a! t+ j* f    for i in range(num_rows):
        for j in range(num_cols):
: W4 M+ ~8 h. P4 u" l            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张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
- M; l* J" e& x& D: E7 l: q
# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
6 H& R0 |4 m, S+ q6 |5 Zimgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)



列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
% |- }; e: a0 p( l" g                [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],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
! t6 n, m! }% }. W0 ]8 ^% x: i* w                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
8 P2 v7 A1 t* B               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
' ?. v6 A  q: u有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
+ V9 S2 m' o3 E
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
' j8 c4 m7 W, l+ Q9 Kfor i, colormap in enumerate(VOC_COLORMAP):
6 ?$ u5 V& [2 z0 ]- D6 @- _    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
9 S7 b/ |9 _0 V7 a0 g: W    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
8 X6 ~4 \4 n3 _
# 构造标签矩阵
7 _( m8 t8 F6 _1 x4 udef voc_label_indices(colormap, colormap2label):
. |% A+ [8 |6 z2 K/ y7 ~    colormap = np.array(colormap.convert("RGB")).astype('int32')
0 q2 l( J; x' c0 w" C    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

8 Z3 e2 |- |6 r- O& G; S可以打印一下结果
! v8 {9 q8 r4 a* A- I
; R5 `4 D3 z9 E  Xy = voc_label_indices(train_labels[0], colormap2label)
: ?; U' h4 _4 |" N$ L8 _3 Bprint(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

9 @$ I) ]; l4 _8 x* z, |" q  n0 Z( W2.4 预处理数据

7 ]7 n0 d" t( w) e6 w+ E, j7 e# f9 L7 u" J在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
6 I1 d5 n0 f% s9 Q
; ]) j- n# B: Q  l$ c  j
% j/ P% n( C# ^7 j2 Jdef voc_rand_crop(feature, label, height, width):
- A7 j- A4 ?( y, a/ V! H    """
9 M$ I4 ?  v' N" d! T4 A6 q    随机裁剪feature(PIL image) 和 label(PIL image).
; G8 |1 H8 @/ B    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
/ U- C' H5 `3 j) W    Get parameters for ``crop`` for a random crop.
7 I: |3 V6 x1 t9 G# ?+ I# V    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
. L0 L- K" n( U3 \# a4 E    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
% Y; g" W6 D7 u    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label
# s0 E2 H6 n* ]; P# B
# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);
  L9 ]8 p* ^1 ?
6 ?3 }2 B7 R: R9 B

: d, h; D8 r. m, ~8 [# z

2 [7 W0 R" z. t9 M6 i% P% G
" B% I0 `! S! `6 K3 自定义数据集类
$ M4 k1 q# {- s; @1 _4 W
3 Y+ ]: E. q( x5 f4 R! ~3.1 数据集类

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

__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
7 s& R) ~1 f5 d3 o* Q" _. }
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
" l$ y2 t' X# D" N        """
+ E+ ^7 @: ^" v. ?: ^( k1 ]        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
4 S0 X- m9 W& g+ B* d; J        self.rgb_std = np.array([0.229, 0.224, 0.225])
        self.tsf = torchvision.transforms.Compose([
. v+ \& r" V) B9 u& n. k            torchvision.transforms.ToTensor(),
; t, v8 C2 _. M; r, \            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)
2 v1 x, l3 `; _' w3 ~6 C7 D; J# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的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')

    def filter(self, imgs):
        return [img for img in imgs if (
/ s! Y9 P8 s+ g" {& d            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
* K) t7 @2 H: a( N1 J: W1 _                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
; u* H9 y) O2 [. W4 @
    def __len__(self):
7 _  Q% y( B- E# Q, ]: H( w7 X: x  Y        return len(self.features)
3.2 读取数据集
* s4 W: `* W' D1 b$ n$ k
( r. f- d+ p5 E& Y' [1 Z" N' f通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
! t; A9 Q5 m8 C( G1 v5 N) U
batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
' d' _" F! I5 z( Q% g* ]3 pcrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
- j+ ^/ Z- E( o. Vmax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
, p4 ]7 d: g1 F, F2 K
# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
( m) k" q% V+ M* Q  J5 r/ u- Tvoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
! V% r( {2 o- k7 R8 A1 @
$ r$ _- O7 k- F" O$ _! ~3 U# H& |# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
* I! p/ F+ ~# z( T7 N; Etrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
7 T, @1 }6 C2 e# t5 R3 K5 t9 n2 A                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
6 C8 P# x  Q: V. ~% w! d1 V& edataloaders = {'train':train_iter, 'val':test_iter}
) I: B8 {1 T5 f! X9 A, b' x$ Jdataset_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')

; a2 O- P1 r0 [* y( w6 n* v8 O2 J4 _num_classes = 21 # 21分类，1个背景，20个物体
9 v3 i8 e8 J$ K0 A4 bmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
' H0 E8 U& y) P' T( }. o$ k  zfor param in model_ft.parameters():
1 z: ~3 G( p4 w8 @4 M9 }* d$ R    param.requires_grad = False
4.2 修改成FCN

; C8 v+ V8 f; a1 ?全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。
: r; g+ x( {4 k: y. @
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
6 y) C5 p+ K9 e0 }3 o4 f8 a* v  S5 y
可以先打印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 倍即可。
* s/ m: W* y0 i) Y3 _- [# W  X
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
( A  v! l2 v; R! y! S( J              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小

# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
* O, q: g" J$ Q* k7 F2 \4 P* s/ Q; bprint(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
2 A/ c$ ]% Q; C$ v* e. g' q9 R
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
/ A5 S; E7 _9 ]; F$ Q+ u6 U; ^+ y#     print(X.dtype, X.shape)
' L, W/ |* H( ]. Y* W, p#     print(Y.dtype, Y.shape)
#     break
. t$ z. O2 b/ Z& W. D

4.3 初始化转置卷积层
$ K: E  b3 S! G0 o2 _0 f* C* J
在图像处理中，我们有时需要将图像放⼤，即上采样（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 k; P& H0 T3 Z3 o( o
( U# Y. J9 g$ l' v( ]" ]1 K' u; h# 双线性插值的上采样，用来初始化转置卷积层的卷积核
; m$ D5 r0 ~! x# M/ E$ e* N* xdef bilinear_kernel(in_channels, out_channels, kernel_size):
; S9 [9 m# k5 O. ]9 p    factor = (kernel_size+1)//2
: K# Q8 M0 m9 r! H  R1 s* E    if kernel_size%2 == 1:
7 @1 e' Z! l( ~  S, r0 U/ s        center = factor-1
! S2 x8 y6 y  m8 x1 s1 r* x    else:
& j. m) [& W3 A- ~0 O" }& {        center = factor-0.5
    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')
    weight[range(in_channels), range(out_channels), :, :] = filt
6 I; U( v4 Q" Q0 q3 j    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight

  B6 b7 F% t5 ?$ Q( _( j
! z" r$ T, O, W. g0 d2 k: {1 t在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
9 p! L+ t% x0 J9 y4 ^' {: amodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
. v6 Y; f* x$ D* m0 a
- Y3 P4 D4 M4 \, U# w
" p1 M5 ]/ F/ `1 g& e
  g3 E" f4 S7 G! t9 {5 训练模型

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

/ i. }0 D9 M9 j9 E+ ?3 S( e, a! i. Gdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
6 y  }  ~$ E6 j! P/ O    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
: M6 w4 V# `- f( k5 f' @( i    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
4 ]3 U6 ?" E2 h4 w        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train()
            else:
( F5 ~# H7 q; g, y! U& N& m3 T* s2 S  R: |                model.eval()
: O+ Z& t2 x2 m3 H! \            runing_loss = 0.0
            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
, h, P' W2 H0 {# N( F! R5 a                optimizer.zero_grad() # 零参数梯度
                                # 前向，只在训练时跟踪参数
- `8 x, ^, T. A2 v' k2 }" ~( ?                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
) ?, e0 Z- O7 Q6 |) K                    # 后向，只在训练阶段进行优化
                    if phase=='train':
                        loss.backward()
4 }  M+ _; c8 c                        optimizer.step()
5 X" H5 k, `1 x6 @8 M8 G' x- }8 V                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
% q/ K% g& _" h' N' M$ S3 P! C: H' g: g
3 H& `0 l- e# g( m; U8 l            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
3 ^8 F! a( N9 |            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
7 t3 s* c. I  J3 b/ v, Z                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
9 ^' {" o+ B; k9 A6 _% k5 o+ s- _        print()
- p/ ^7 Z& `; g    time_elapsed = time.time() - since;
; ]* `3 S, {3 Z    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
! l. S3 l$ }3 R( [) n    # 加载最佳模型权重
+ C' L* D* Y: `, Q" }    model.load_state_dict(best_model_wts)
! C3 `2 S( G6 ~3 U- d% C  m7 Q    return model

3 O; |- q" ^; d% J' ~, z9 G2 I8 y下面定义train_model要用到的参数，开始训练

epochs = 5 # 训练5个epoch
1 k0 M( C7 Y& F8 a2 N: Ocriteon = nn.CrossEntropyLoss()
( i" V" j8 t4 Q( R7 p+ B' yoptimizer = 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)

# 开始训练
, \; I4 X# [; B& z; f' W! A9 G; Nmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
& U/ Y4 N& {: \5 C6 o
6 测试模型

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

def label2image(pred):
$ g7 ~# u& j+ Z; H; ~+ _: i1 s    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
' x/ [+ ]; ]9 j0 T. {: C: X7 ]    x = pred.long()
6 l$ y  c' B& f4 _    return (colormap[x,:]).data.cpu().numpy()

1 a) i% X, H% o0 Q

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
# j' W" N5 A- I* H0 C" Nstd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
' M7 D- ]5 j7 V' R4 r5 i5 Udef visualize_model(model:nn.Module, num_images=4):
6 n  V# W' T) U( K( Z    was_training = model.training
0 k( f  P% h; Q- g% s0 j, t, _2 O    model.eval()
9 c3 O: r+ u3 ?& Y5 X    images_so_far = 0
* q- V8 x3 z8 p6 j    n, imgs = num_images, []
5 L- c. M  e9 P4 m    with torch.no_grad():
3 F% M' |% A, o5 [        for i, (inputs, labels) in enumerate(dataloaders['val']):
! v9 F2 j3 S! u+ C5 D            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
7 g% p- F5 s* s$ B& c. |7 ^5 \            outputs = model(inputs)
            pred = torch.argmax(outputs, dim=1) # [b,320,480]
4 g7 [5 Y5 W4 t, w' ?3 c            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦

7 E" ^, \2 ~# A& h. z2 l6 O/ R            for j in range(num_images):
                images_so_far += 1
                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
; N. @1 j+ O2 T2 `) P8 [2 y1 O                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
9 w- y" D6 ]! V                if images_so_far == num_images:
                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
- J$ N/ @7 j( h0 z2 m7 _9 |4 V) u                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
- x# R; r1 N: l8 }7 [) g                    return model.train(mode=was_training)

, y* w3 s2 R# L. C1 E+ I$ b# 开始验证
# r4 k# J' u& k! R3 \% r! lvisualize_model(model_ft)
8 @8 V' c* U9 _2 i9 I% O! r6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
7 s/ i4 I- \& u- p6 zdef predict(img, model):
% n; P* G' D5 o- q% B" b) r3 m    tsf = transforms.Compose([
3 K/ F  R8 X4 r6 I            transforms.ToTensor(), # 好像会自动转换channel
3 K2 P9 \4 F' Q* F( `; L& `            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)
: Z1 ^, N. ?' s8 g    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
6 u& r% u! o* Y1 F' D* ?* s3 R
def evaluate(model:nn.Module):
8 {- A9 V* E  n: Q8 r" \    model.eval()
. W( G' G' N0 P    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
    n, imgs = 4, []
& @5 Z- g( u) n! z    for i in range(n):
1 O9 [9 Z2 r! L; g; g$ z        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
4 w! }) ?4 F) Z8 O6 ~2 ]        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
5 ^; h3 H$ _2 V% p: M$ e1 H    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
" N% k) Z& }2 g- [( M. V8 ?
# 开始测试
evaluate(model_ft)
8 v3 a9 F+ i5 E) ?7 s
7 结语

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

Epoch 0/2
8 H+ ]  \; T4 U$ _7 u4 ]----------
$ K. h' v4 H! p& g# i  ttrain Loss: 1.7844 Acc: 0.5835
9 A, Y7 @# ?' ^  c" }& P' N, zval Loss: 1.1669 Acc: 0.6456
9 T, p; j: ]! O- r" U) ~
Epoch 1/2
----------
train Loss: 1.1288 Acc: 0.6535
- j# R+ H' _6 b7 Z$ _3 Kval Loss: 0.9012 Acc: 0.6929

8 W# ~7 p% |& Z1 a2 Z, AEpoch 2/2
, r; a8 z4 a5 [" [% x! m----------
( Z7 c) {" V( v# k  Ctrain Loss: 0.9578 Acc: 0.6706
. R  |. o7 W$ g* B0 n- _- C& yval Loss: 0.8088 Acc: 0.6948

" H; i6 y" K/ [! b# p  _; I6 S" iTraining complete in 6m 37s

( R3 G1 C& n. E$ @
. s  P$ q1 ^7 J

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

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

0 L- k% f; R: X语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
( N7 Z, q# V: g% ~* o, l. {1 I
  J. K7 D0 P' @3 q3 i# s. ODeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
" I, k  X1 B  _* _GCN 通过全局卷积网络改进语义分割[论文]
, l$ C# w- f2 M5 m" y; VUperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
! V9 P0 {$ g3 d8 p: Z1 dU-Net 用于生物医学图像分割的卷积网络
/ f% f6 g+ a3 }& E; OSegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

8 I+ t1 Z% C5 a2 V常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

对于损失函数，除了交叉熵误差，也可以用这些：
! r* H2 W  B/ E/ d
4 F) ?) w2 C$ P& d- vDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
0 p- {+ @, s' S6 A4 [4 }2 e




————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507



, _; i+ J1 L; s- \/ z5 ^