Pytorch实战语义分割（VOC2012）

杨利霞

# H/ r) p6 F  Q, h* ^  x
7 J1 N& F% \$ i  VPytorch实战语义分割（VOC2012）
) h4 _- v% j! x0 y本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
1 m" z/ r: P9 y8 z2 I

语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
2.1 导入模块
7 v6 P& R3 ~: C! N; n* a2.2 下载数据集
% f% N" T) n1 k2.3 可视化数据
& \" W* d6 K7 U! I% `% w* f2.4 预处理数据
, r! [; f0 R. Q$ c3 自定义数据集类
3.1 数据集类
: o8 M# x9 Z/ x( K* [; n$ f3.2 读取数据集
7 k. u! R2 `! Y6 X$ Q4 b$ H4 构造模型
. \; M4 b; y+ j$ ], `: J8 Q4.1 预训练模型
4.2 修改成FCN
6 y+ p/ Q* u0 q7 Z. N/ P4.3 初始化转置卷积层
$ g2 S8 r! U- v- f( b. d5 训练模型
/ n$ ]! _. Y) {5 n& y7 l' Y6 测试模型
) a0 x% m& _4 b  z+ o5 g+ i( q6.1 通用型
! F  Q- I0 g8 g+ C6.2 不通用
7 结语
0 f! Y/ t% u9 X4 D- y1 E1 图像分割和实例分割

! }- ~- i) b5 p5 r2 l  e计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
; F3 m& c/ t% l" `; M) B
图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
7 O; l1 j  B) ~  j' O8 O) x$ j实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
" ]0 k+ x4 d7 h: z
2 Pascal VOC2012语义分割数据集

2.1 导入模块
- P0 F& b! e. Vimport time
$ U% {/ b2 s1 [, W- H+ I5 l8 Zimport copy
5 }: b( z) p* R8 L+ a& \6 ^import torch
6 q6 o" q5 e! |7 \9 mfrom torch import optim, nn
% y" l' P/ t! N8 D- M& Ximport torch.nn.functional as F
' u; h! M8 u) n  o  ^% w7 x3 x& bimport torchvision
# c$ w: L, }  B% Y! I. p" M% v# zfrom torchvision import transforms
from torchvision.models import resnet18
" i1 B  v$ X0 e# qimport numpy as np
0 K* y$ ~: ~2 r9 bfrom matplotlib import pyplot as plt
, H% }& ~8 I& e" bfrom PIL import Image
import sys
" M/ R$ E' p! W6 Ssys.path.append("..")
from IPython import display
  r5 Y4 Q2 e; F. W( Ufrom tqdm import tqdm
import warnings
7 u+ s- \- E8 K) z. ywarnings.filterwarnings("ignore")

2.2 下载数据集

) T9 g- t" t  [# W. x语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
" ], D- \* m+ B5 a( M8 l- _
! z; A% A' W  q/ O" U% J
# l8 `3 h4 ~1 N2 G$ HImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
: y7 O# m1 l9 G2.3 可视化数据

6 ^" n. D/ g+ A6 Q6 m: y5 \定义read_voc_images函数将输入图像和标签读进内存。
/ P: p% K( \+ H) I0 j
def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
) G( M# _- z& \' o0 X1 }    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
    with open(txt_fname, 'r') as f:
% w+ E9 }9 \) G7 I* T9 ]1 g        images = f.read().split() # 拆分成一个个名字组成list
    if max_num is not None:
3 ~# i$ R5 G$ o/ D        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
5 I* y" `0 z5 C8 V    for i, fname in tqdm(enumerate(images)):
3 N: q8 p4 J8 P7 a        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255

" L6 I' _- M9 ]* k* R2 T$ |定义可视化数据集的函数show_images
# J* `% G; n( o  H9 A* r
( R5 p/ P6 L/ ]5 f3 K+ @" z# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
( v, q, W7 l& }1 X2 d    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
$ M7 y8 ?0 d7 I    # a_img = np.asarray(imgs)
. h5 u: d3 g! i8 J7 A* ?5 G    figsize = (num_cols * scale, num_rows * scale)
    _, 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])
            axes[j].axes.get_xaxis().set_visible(False)
; V7 I7 E7 Q7 ^' Q6 \# T            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
1 v! l( B% L! B0 g3 \3 Z
定义可视化数据集的函数show_images。
1 l5 T: J' ~& j8 H7 J0 }9 s5 I" {
# 这个函数可以不需要
- ]# [& K1 ^5 h( @( G. `9 Vdef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
- e4 O" z! \7 m/ D% X    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
  q6 \, F7 C. }' u. @: I! i7 R    plt.rcParams['figure.figsize'] = figsize
. ?  q4 f1 e4 m- ^/ c( u; ]
def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
( Q5 ?* d3 [1 e    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
        for j in range(num_cols):
' `* I- K: b: |& N9 Y: Y- ?& B) s            axes[j].imshow(imgs[i * num_cols + j])
* p1 A! p5 p/ |8 W            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
! ?) s: o) D$ C& ^8 Q( |    return axes
( ~) w$ }  L' j. M1 K! c画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
, N7 w: R/ W$ ], w3 b8 F
7 ?3 Z5 w% V, O/ }7 b6 X  Z# 根据自己存放数据集的路径修改voc_dir
7 m7 Z+ O" U* {5 T, {voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
* ^: M* V7 f0 T7 B+ c  r! bn = 5 # 展示几张图像
! g- Q& O: Q" C2 m: wimgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
+ l; A6 a: |$ ]6 n  n7 b) w. T
: B! i9 q6 i% O) R' R
$ G" j: J2 x" V
列出标签中每个RGB颜色的值及其标注的类别。
# c+ _$ l' L6 E* D5 k4 |2 X' C# 标签中每个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],
$ }# V/ [3 D# S5 b# N  U# X                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
# v- q: K. f; T8 M" x                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
# 标签其标注的类别
, }7 i5 a/ S& B! W% ?" S4 _. iVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
& Z. \0 V9 u6 U% @2 |$ V0 @- ~               'diningtable', 'dog', 'horse', 'motorbike', 'person',
' X: G9 o3 t9 ?. R! i               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
" e7 o+ ~( V8 d8 G1 t; R有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
. E, A$ p, B& l3 c. k- h; Z& fcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
4 i! O# m0 Y6 Hfor i, colormap in enumerate(VOC_COLORMAP):
9 I: [9 x) I* T. G    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

, X* d( a% B( }+ B/ t/ f# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
/ k- u2 E& ~4 K& e- b    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

* m/ Q# k; j3 O4 h: f9 q0 e  j% Q% h可以打印一下结果
! H8 Z  d5 J- h1 P  ~7 J
& n" u  Y! l0 _4 A/ w; Yy = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
6 i# ^3 b0 [5 [* S
2.4 预处理数据

, m% T, i7 `4 v# c# ^, C在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
- o, Q( \. Y" d  F5 o" t( W

) u" b3 d1 V: ?* g8 Odef voc_rand_crop(feature, label, height, width):
$ T8 Q8 n8 D7 D7 i3 Y7 K& r    """
) K& p! a, o7 S. a; |( s    随机裁剪feature(PIL image) 和 label(PIL image).
2 m6 x' v, @6 N3 A$ C% o# M    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
7 k5 I, A0 ]1 T. I) `% b- m        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.
5 @' Q1 m& L/ l& N+ P    """
+ W* T" }- z  s8 u# f& r/ X    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
9 ]% ]* W  X: b+ X6 N    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
+ |7 _9 ~# E6 L7 b    return feature, label

# 显示n张随机裁剪的图像和标签，前面的n是5
& j0 ]3 E# w& v# M& @/ Yimgs = []
! l- S1 M5 ?6 S& ?; s( Ifor _ in range(n):
# {9 J4 D% J, ~6 @% F    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
( F' c& \) o' L3 H1 l7 d- Vshow_images(imgs[::2] + imgs[1::2], 2, n);
( a* y) n& i' C, I2 S
+ C- W7 t6 e5 x+ T6 W( I
5 X8 @5 m" K# m* b+ R% h* ~

) K. F/ R1 K) A- N: t. V" B5 V

3 自定义数据集类
2 D+ m1 S+ y& T) B& C
1 E. C! y* @: l% |) N" x# \# x3.1 数据集类

8 ^5 s3 [5 r3 R: s" t( \" qtorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
. v% s6 M: p# [3 t5 G4 \4 D
3 k) F$ x* r/ X7 ^' I, Y__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
5 z& _) I: r" g- R由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

+ O5 r6 A" a/ r9 D5 h6 oclass VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
3 i& p# u- M' R7 u0 Z        """
' U0 T5 V: g: j  m1 ~        crop_size: (h, w)
2 P' ^5 B6 u5 H' R$ S3 g5 n        """
        # 对输入图像的RGB三个通道的值分别做标准化
7 m. q" z  F) d0 |8 `$ j% y$ ^0 w7 h        self.rgb_mean = np.array([0.485, 0.456, 0.406])
: a$ `" v2 R3 U; ~  E        self.rgb_std = np.array([0.229, 0.224, 0.225])
        self.tsf = torchvision.transforms.Compose([
7 O  i+ T& ]- q7 V1 g+ g; F            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
: e) T4 }/ h9 N0 ~3 k' D        self.crop_size = crop_size # (h, w)
; I  V% l$ u% j        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
8 [: }5 U1 |# [9 f# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
$ t0 {( M( S$ d' u$ V        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
! Q, U/ S; J) L3 h; P/ R/ Y5 R9 o        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

    def filter(self, imgs):
        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

    def __getitem__(self, idx):
6 t& C3 j8 \. E/ N3 H" X( \        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
                                # float32 tensor           uint8 tensor (b,h,w)
* i. {  B% f* |        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
% D, U. B/ X$ ^$ |5 @  M
    def __len__(self):
        return len(self.features)
3.2 读取数据集

通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
1 o3 b" l1 J% W8 h" B0 Z1 L! y* b
batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
! b% T& R& o2 y$ w. {* rcrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

4 u/ j+ T; t0 G$ w# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

4 ~. R, Y7 u# N8 m# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
9 ?: h9 v) L2 Znum_workers = 0 if sys.platform.startswith('win32') else 4
train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
, I& N/ w) Y' z, d9 }                              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里
dataloaders = {'train':train_iter, 'val':test_iter}
dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

0 j! O( e" @! w9 g: O! I4 构造模型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，表明要加载使用训练好的参数

  W3 O- h# a" W; T# 特征提取器
: ]3 v) g6 W& H6 ufor param in model_ft.parameters():
6 |0 i: H% ?7 V7 U    param.requires_grad = False
: e9 W/ s( x! H( I  w/ A4 f4.2 修改成FCN

全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
- e: ]5 j' I! x4 v, v" Z9 a5 [7 {9 j最终输出的通道包含了该空间位置像素的类别预测。
; H4 x/ E& X# N" s7 n
- _) h0 I) E, h+ t  t对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

2 B1 V8 L1 U6 y7 B9 m可以先打印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- |3 S
" Y7 n4 S& a+ K- Q/ ~' _: zmodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
) s# z, C: _$ l  V9 x3 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做一个测试
# e  A3 J6 {! s- |x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
3 @; C0 q. V% l/ x7 ~
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
6 R3 l7 V- l: H6 u, u#     print(X.dtype, X.shape)
) u5 K' [8 J6 C) I7 H7 S#     print(Y.dtype, Y.shape)
#     break

  A- f3 D2 f& T, C
4.3 初始化转置卷积层
5 x' I9 n- C/ n
; |  Z$ y9 z$ |% o& c' q在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
% z, g8 ~: _% M8 q" T/ V  a$ d  t在坐标 (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函数构造的卷积核的转置卷积层来实现。

) V  O4 i# i! Y8 O! H' g% d0 v& t
# 双线性插值的上采样，用来初始化转置卷积层的卷积核
! \9 Q* ~* h" T6 l6 Sdef bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
+ ?! r9 @7 z0 t' W3 \    if kernel_size%2 == 1:
# M. y/ x; p! n  E7 F. _        center = factor-1
  U! g) S! h* t) D8 t4 `9 x9 D' Q. S    else:
: e, x, N6 Z' w        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
" X) [- g6 h0 x) m4 `    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')
9 U( o9 x2 R0 p    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
% C- d4 n# B# v6 ]; Q" n    return weight
% F2 Q0 s+ u' |0 _1 O5 J
8 r. m! o) ~0 D5 `
$ q6 f" [" o. D1 N+ o4 y# @7 _6 x在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
1 N2 k" s' }. x4 l' T# F7 ?! wmodel_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
! S8 v+ v& N) @' a5 B2 _

4 b8 }7 G  }; p* M+ {
5 训练模型

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

. G  h& b4 j- _; jdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
2 L  C+ n' U( n0 J    best_model_wts = copy.deepcopy(model.state_dict())
- c# @- ?  f7 f0 b! C9 k' a" I    best_acc = 0.0
: |$ r7 q; M  C8 o    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
        for phase in ['train', 'val']:
            if phase == 'train':
% _, I# u  V4 C1 y& u' j$ |                scheduler.step()
$ C  `% x# O2 B- t8 K) ~+ O                model.train()
8 X4 Z+ Q2 Z/ f3 O; d. i9 l2 h            else:
                model.eval()
# T  [* J. s9 G, A# s; G            runing_loss = 0.0
            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
                                # 前向，只在训练时跟踪参数
7 v8 r, l7 h) j, U6 u                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
/ n2 y8 L1 U8 p. y                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
6 s; d3 Y: _( p! W: f9 ~                    if phase=='train':
* f+ r( X' n- {                        loss.backward()
                        optimizer.step()
- P. S, @9 @- ?                                # 统计loss和correct
5 H0 ~7 O1 T( o                runing_loss += loss.item()*inputs.size(0)
3 J. G( y% T8 A: V* x4 ^                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
  d+ u2 ~8 U4 ~, L8 m
) d7 e5 d4 W- y            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
0 Q: Y* a. Q  `            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
" H" V$ ~$ S; r6 K5 R# a' A, t            if phase=='val' and epoch_acc>best_acc:
  n( `+ v2 Q  w, Z6 K  z                best_acc = epoch_acc
: s+ R8 X$ W2 ?$ b                best_model_wts = copy.deepcopy(model.state_dict())
        print()
    time_elapsed = time.time() - since;
( k8 [1 @+ E  A4 y$ v    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
3 ]) t. [3 w& B    # 加载最佳模型权重
4 F: l6 W9 o# n% j$ w  W8 D) Y    model.load_state_dict(best_model_wts)
% @1 W2 \) b# ?$ N    return model

下面定义train_model要用到的参数，开始训练
, Q! o& Q1 q$ G. u  Q1 l
epochs = 5 # 训练5个epoch
- H* [( a3 e, |# c" U5 `criteon = nn.CrossEntropyLoss()
1 d( w5 J& o. o" b# I6 @) Soptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
; {" o6 Y& T" Dexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

6 P+ U8 a3 n( T" b) {# 开始训练
! x! y) h! H/ Z; t; ]' a9 amodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
6 _) n! V1 W3 n" k! q
$ l  f. d/ \5 c. R2 F; E! Q6 测试模型

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

def label2image(pred):
    # pred: [320,480]
3 f1 x* _9 E) ?    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
8 G: b; p! c3 A7 l( z    x = pred.long()
4 r7 @4 ?) P' Z: p( U$ f    return (colormap[x,:]).data.cpu().numpy()
. N1 J3 r# `, M

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
: s9 i- q) {0 G4 ustd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
6 i& x1 Z$ ^: [8 O6 h! Qdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
: A" d/ t. [6 n2 `) A7 Q- H8 y' [    model.eval()
" X  P0 R; R. _    images_so_far = 0
4 N, ^5 t% h4 ], L+ l4 f0 B    n, imgs = num_images, []
    with torch.no_grad():
  h0 K8 d7 I2 V  G        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
            outputs = model(inputs)
0 A* Q' t* ?8 [% b4 G% c            pred = torch.argmax(outputs, dim=1) # [b,320,480]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦

            for j in range(num_images):
                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)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
; I+ x8 I6 K9 {                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
/ l+ V. {+ h+ V) }2 m% U                    return model.train(mode=was_training)
# M+ D6 P" j$ a: X, v
# 开始验证
visualize_model(model_ft)
) R8 R- a) Z4 i3 f& f, k6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
) `4 w* s5 T! E3 x  c: Q            transforms.ToTensor(), # 好像会自动转换channel
5 A: E4 w2 m% A7 y' s            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
/ ?8 v2 S3 v1 I' i: b8 z" W6 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)

def evaluate(model:nn.Module):
    model.eval()
' g) o3 [, Z" s$ }, r! a# o    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
    n, imgs = 4, []
. ?; E  ~0 `: ^2 Z/ u( `; q    for i in range(n):
8 `4 @. o0 ?- C6 p6 d        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
/ T& w9 O6 o  W( G* C8 E; o9 `        pred = label2image(predict(xi, model))
! E. ~2 B7 g9 m, z( J        imgs += [xi, pred, yi]
" H4 ]$ X  u/ z5 S2 i    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

# 开始测试
5 ^# K: r+ t. t- jevaluate(model_ft)

  [5 y; m% t6 v5 W! t, y1 t7 结语

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

, X2 h' ]3 ~( A' jEpoch 0/2
----------
6 f5 q" I3 L. \train Loss: 1.7844 Acc: 0.5835
! R/ V/ Q. \' K+ m- F1 H6 L8 Xval Loss: 1.1669 Acc: 0.6456

, O' W, G; m3 ?* C3 b% j2 qEpoch 1/2
' J2 ?8 S& D. z----------
, \& T4 g1 A9 U  Z9 t; K. m. |: [train Loss: 1.1288 Acc: 0.6535
5 ^' j9 V- f: F% _0 Rval Loss: 0.9012 Acc: 0.6929

Epoch 2/2
3 a5 x0 S' b& T! E/ C& S7 @9 \+ C----------
( E4 `! w& q! S0 X& ^train Loss: 0.9578 Acc: 0.6706
4 G& P6 N+ c1 \  c6 ^$ \3 T; j# H5 \val Loss: 0.8088 Acc: 0.6948

+ C4 W* G/ k& V7 C/ p. {, aTraining complete in 6m 37s


2 P$ [7 r8 v/ A4 U& ^& t

当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
% t+ `3 U3 M( T2 J
对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
& w! y' I- J7 x% ^
" O4 A4 r8 T1 n
语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
! I# ?1 [- X: S  A: p+ ]0 n# y
' I  V3 i$ Z( i0 t* bDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
- u+ K: |2 [: V- O( p: I5 f8 {UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
% s, f. {# c& |! C$ H+ t+ UU-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
0 I, B) N3 N/ @; n6 a! R还有（DUC，HDC）、PSPNet等。

常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

5 M* M( ?2 l; o& T对于损失函数，除了交叉熵误差，也可以用这些：

Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
* ?$ F& N: ]/ [* B4 lCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
& `2 b. P$ m; x) ?+ v1 [Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
* B$ l! E) x$ S) a! `1 sLovasz Softmax 查看论文：Lovasz - softmax损失。





————————————————
4 e+ l* f6 s: K: o/ m& F版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
: D# q7 n; h  ^- ~, R" O0 k