Pytorch实战语义分割（VOC2012）

杨利霞

; S" q0 u0 l  N: ]; K+ y
" K$ D% F/ U6 G( w0 u: MPytorch实战语义分割（VOC2012）
+ a- V* ?( L- Z* k9 f本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
/ x; M) Y) A) h) S' x. k; s语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
, c( Q5 V) b* N& }: G# g5 _) D; _
! |7 s( P, t$ ?  K( L" s* Q1 \
语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
2.1 导入模块
2.2 下载数据集
& L; ~; t: G7 h/ u0 ^2.3 可视化数据
2.4 预处理数据
' w# Y2 `7 f2 B3 自定义数据集类
3.1 数据集类
% d4 |: s6 B2 \2 S3.2 读取数据集
4 构造模型
4.1 预训练模型
4.2 修改成FCN
4 @4 \, j& p' k$ G! P4 I0 L4.3 初始化转置卷积层
' E" l( [  c( Q0 l. g5 训练模型
6 测试模型
6.1 通用型
6.2 不通用
7 结语
1 图像分割和实例分割
9 P& [: ]( d: F" \  u' t. @; s7 p
0 [- Y! Q+ r7 e4 @4 Q( e计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

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

7 G$ s! }/ V5 @1 N! ]: D4 z3 R2 Pascal VOC2012语义分割数据集
8 B. F3 Y9 y& J! g
; ]5 n+ O0 J, S2.1 导入模块
import time
0 t/ ]5 E7 s+ t8 l( r' Q* Iimport copy
0 `, w1 F" U+ ?1 Y; z4 U7 W( c! D6 @8 Ximport torch
from torch import optim, nn
import torch.nn.functional as F
3 A- T& I8 e; N/ Bimport torchvision
from torchvision import transforms
from torchvision.models import resnet18
. o5 O. N) L. Q  a. A! gimport numpy as np
from matplotlib import pyplot as plt
# D+ R8 m* f% ^& A; Gfrom PIL import Image
import sys
sys.path.append("..")
$ l" B* o$ U$ x( c* u3 Jfrom IPython import display
from tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

5 N& I1 X6 i$ J% u9 A- n2.2 下载数据集
! F7 C7 u# `+ s7 W) \/ H
, H4 t7 W0 P$ s5 |, @) R! m1 [语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
" Z, C; {! C* Q" R
' Y" W6 B" L9 S* X: p( q6 N
ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据

! n0 v" {. c7 s定义read_voc_images函数将输入图像和标签读进内存。
* Y6 V/ O) D1 Y7 E4 k
def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
" W8 k5 s2 B$ y/ c) a    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
    if max_num is not None:
        images = images[:min(max_num, len(images))]
2 K, c! s# Z; F" C* i    features, labels = [None] * len(images), [None] * len(images)
% o$ R+ ~% f2 M7 Z* V$ h% d+ n    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
. P0 @. L0 d5 [" l, ]* m# _# }        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
' {5 G* {0 N# z( x. k        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
1 d& u) g% a1 u7 A+ h6 b  ^    return features, labels # PIL image 0-255
: F% R9 Z' w3 n( {. Z* }
定义可视化数据集的函数show_images
; j: w/ i* F4 ~3 L/ X$ ^( o  E" {6 i
- p$ D: {  R3 }! C+ s) n, C# 这个函数可以不需要
: t; H5 l4 `' ^: }6 Ydef set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
$ ^7 w' s, C# N0 O% }/ Q    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
6 K9 b/ b2 i( r, V" Z
% ~4 S# k1 Y* V8 P6 P" H- _, Odef show_images(imgs, num_rows, num_cols, scale=2):
" |2 S2 p. d# x+ L, I, ?9 p- A% F    # a_img = np.asarray(imgs)
! O1 @- o% W: j, `8 ~    figsize = (num_cols * scale, num_rows * scale)
2 c+ t6 N: T7 x, c- t    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
7 r$ J2 w& G! A    for i in range(num_rows):
4 s) f( ?& b$ I+ S& C        for j in range(num_cols):
3 \% F6 X7 I3 e0 H% c4 X! t4 N            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

定义可视化数据集的函数show_images。

% M: J1 w& }' P* @: L) i9 r# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
7 y& A7 ]" T9 Q7 @% f; V    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
3 D4 s9 M5 e% s1 U    figsize = (num_cols * scale, num_rows * scale)
+ y1 i, t* x4 m' z2 f    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
% i4 A, V* P0 ~) c    for i in range(num_rows):
        for j in range(num_cols):
6 ^3 Q/ W/ w& ~% k            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
' Q& o9 |: o! j4 L; C  M. j            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
1 v6 |2 V" Q- O. ?; G    return axes
6 U9 A" s% l5 W9 U画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

+ z! ^/ l8 ^3 d- W7 u# 根据自己存放数据集的路径修改voc_dir
4 Y2 n9 H  w% e/ |voc_dir = r"[local]\VOCdevkit\VOC2012"
/ B4 B- W' B# ?# Ytrain_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
% B. ?. n" g& H& T* y. V3 Wimgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
5 H% q3 v# N2 ]" ^6 _2 C# M- T


3 f- q. l" Y- D; o$ n: X列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
% x7 ?& G6 ~, S                [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],
( m+ u" W! Y8 |. \/ K                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
! d" R0 n% e/ t4 O1 Q3 R( g                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
9 P1 W* q6 |2 d               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
  V: K& F5 j8 e" v! S2 }  Z$ ^8 h, P               'diningtable', 'dog', 'horse', 'motorbike', 'person',
( n' R5 @+ o, M) ^( i" u$ u$ F               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
. x$ A1 V- Z  n
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，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):
& y" h0 E) \# ~1 g9 _" K6 B    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
8 m6 P& l1 c) U) [  _! f! n
$ }  C* B+ Z! {6 G( K8 u# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
6 i- P1 S3 _0 c, F    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
  O7 [& V5 q' h! Y    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
& v) n" _0 w/ U  E0 Z
1 q4 R3 d1 c* H' {3 v( d% f可以打印一下结果

y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
$ Q& g, _: U2 ~/ s& S6 ?
3 t3 T5 s' h5 ^! m* E2.4 预处理数据
$ W( W% u( F  I% e. r
2 d, |& D' r+ C5 m4 V9 t& \在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
& c- ]$ c/ T8 X, S9 b

; L# y. O1 q4 }( Y3 }def voc_rand_crop(feature, label, height, width):
    """
    随机裁剪feature(PIL image) 和 label(PIL image).
0 K5 F0 q* x3 A: p3 b    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
2 Q$ p6 k5 Z( Q4 X# R7 z+ B& n    Get parameters for ``crop`` for a random crop.
    Args:
        img (PIL Image): Image to be cropped.
- Y! C. q: p* w. B        output_size (tuple): Expected output size of the crop.
5 o. ^8 c. T' A4 f    Returns:
' D2 b/ S8 i8 V- ]+ P& p& k2 R; ^        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
# C$ o' S8 F1 M! d4 L8 u& T    """
: d$ _5 ^! s1 c8 ^, e    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
! v* Z; u- u, u* _
# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
/ m& _7 A: H9 t# }3 e- k    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);
2 f8 Y: e  V& H1 N. v- `) {: b



7 \) m+ [6 n# y. G
" d: j) N: `  ?2 C8 i6 v4 t
3 自定义数据集类

3 H. A- l! e+ d! L* M3.1 数据集类
; O. y) e7 Q; e& B
0 l+ c2 _- I$ X7 b; qtorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
) }! f( R* Y* q% \, z由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
        """
5 s0 C& W4 X& N: k' I        # 对输入图像的RGB三个通道的值分别做标准化
) y0 a9 w  J( E4 j, w! ?8 D        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
        self.tsf = torchvision.transforms.Compose([
2 o. ~0 |! h( }) ]4 g+ W) i- w            torchvision.transforms.ToTensor(),
: X1 q! X+ j! I; c6 K            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
& P! F7 n% G7 V' u3 W6 g        self.crop_size = crop_size # (h, w)
# @# q5 B) p! b1 G! l- r: z        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
% _4 t. B2 _4 [0 e' y  t- R# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的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 (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

    def __getitem__(self, idx):
1 L% U6 F/ }$ V5 V- b9 I5 D        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
                                # float32 tensor           uint8 tensor (b,h,w)
3 ?  ^. Z# w1 M- \        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

, f( o! V6 h! X5 ^1 f6 C    def __len__(self):
4 ~' n1 r0 G; v* r        return len(self.features)
1 v- j4 c, h! r9 Z/ s3.2 读取数据集
: G8 x  O& V3 L& p/ M3 I7 ^
0 A" p" u" U  p0 T通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
) y, E, e8 `2 O% B/ a2 X3 p
batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
" T0 ]6 I$ b: imax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

. K2 k( L3 L$ h8 L9 v: d0 V/ N& L# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
  E/ d( d; [# A+ W9 G8 \3 {num_workers = 0 if sys.platform.startswith('win32') else 4
5 P, G% Y# n& Q% ~4 W# E* k. Jtrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
: L* J* y* J7 K* L3 W- R: o/ w6 n                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
( s+ d  ?  l+ v# t" x. i  h                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
2 S* K9 t* `9 f  ~dataloaders = {'train':train_iter, 'val':test_iter}
0 `! Z: l0 ?) m8 Rdataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}
4 S! O7 N7 _) a# r; h" u' n
' @1 X4 S7 {- P! N4 构造模型4.1 预训练模型

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

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

num_classes = 21 # 21分类，1个背景，20个物体
  G: W: [/ A0 Z! ~# `! Dmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
6 }0 @3 i0 L+ S( e+ M$ r
# 特征提取器
& R( Y* f7 _2 ^' j# d  \" efor param in model_ft.parameters():
/ H3 g( [+ y, V0 e0 i0 |    param.requires_grad = False
4.2 修改成FCN

) Z1 n* U5 Q$ R1 ~全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
4 L- U4 Q& r" o最终输出的通道包含了该空间位置像素的类别预测。
8 U9 J; N4 H. n) s
对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

7 N1 P- G/ g2 z; A1 P# C. Q1 x可以先打印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 倍即可。
" C5 s/ \; C! u5 C) ]9 [
7 q; X1 E. t0 \) A& {! B: k1 `; rmodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
4 G4 X" C7 f" Y# d' t* ^% j  I              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
/ V5 `2 T4 T0 H              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
  h& f! V2 H; p) l5 v
* S2 U( I. a8 b) o# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
- j" k2 k" y( l& }1 l& |
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
8 w1 X( j7 R  Q# G, h: h/ O4 V#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
, H" g2 |# r9 W+ W#     break


4.3 初始化转置卷积层
( e& D! j$ Z  q/ \
在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
) X* a8 B# A- R* T: @+ n- m在坐标 (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函数构造的卷积核的转置卷积层来实现。


# 双线性插值的上采样，用来初始化转置卷积层的卷积核
0 G; p/ K2 m/ Z; z5 [2 s8 f  G- Udef bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
3 s4 O( T& O; s/ R7 G" [        center = factor-1
    else:
0 A" b' ~' X! D9 L        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
3 p9 R( t* T, v4 V, j    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
$ [1 W0 Z4 q8 _; z3 b( |- N    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
1 f+ h$ @# \' g8 w. z    weight.requires_grad = True
    return weight

$ G9 r, r' X- E/ Y
在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
" L1 M  K# V4 |- m. M) l8 P6 Z6 \model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)


( ^; `- S9 E& J% x; U
5 训练模型

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

# ?+ D& c7 L# l) k' z4 c/ hdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
  C8 x/ b( n; l8 ]  S% T  h7 ^+ _    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
. ~: X1 ?. |4 U( H% m9 a    # 每个epoch都有一个训练和验证阶段
8 }% s2 p% f7 d4 H    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
7 X# @0 ]7 k  W1 Z        print('-'*10)
+ P9 A( \& ^' x        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
+ |' k: o! r. s1 i( k                model.train()
( y0 T4 h5 C) }8 `% f            else:
                model.eval()
/ d1 U) C7 m0 U7 _' p9 _            runing_loss = 0.0
            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
8 Y, M, |1 \4 ~! u, ]  \9 F( y* v                optimizer.zero_grad() # 零参数梯度
/ \8 Y* U3 L( y                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
9 ]" v6 S! W* @% x; U                    logits = model(inputs)  # [5, 21, 320, 480]
- ], r& E! u- N. W/ L& U& ]  u& P$ v                    loss = criteon(logits, labels.long())
/ X3 }- L6 S' P- S: s- @; |                    # 后向，只在训练阶段进行优化
" y5 S, G: K! h! p6 H- K                    if phase=='train':
2 ^8 f+ Q) n+ s  R; V4 u; ]                        loss.backward()
                        optimizer.step()
. H) y+ V5 I- H! l                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
6 {' u8 v3 [: {) ~( E                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
9 t' W' d+ b: X  h0 e
. \7 c/ e2 ]( L  Z- `7 \            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
  r  _& ^! t: p" s1 z            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
: _4 e% ^; R* l' U# C4 R0 r) C                        # 深度复制model参数
& N8 q8 N, ]; X2 d8 a. x( V8 M            if phase=='val' and epoch_acc>best_acc:
0 O$ ~7 o% V- [* t) x/ S                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
        print()
, H5 f. }/ u8 {- g7 ~5 o: h    time_elapsed = time.time() - since;
. a% }4 `$ f, W( t( F/ X/ N    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
2 O6 Z9 K, Z6 U$ Q1 y* Z    # 加载最佳模型权重
% O- W, l4 ?7 B- K. v# ]    model.load_state_dict(best_model_wts)
5 _/ O# }+ B- S7 ^. t) q! ?    return model

下面定义train_model要用到的参数，开始训练
: e2 J3 o5 B* r: q4 J7 c
; ?$ t% ^7 ?5 X6 Q# e0 t  i  Repochs = 5 # 训练5个epoch
* l+ F$ O0 ~, k4 o* `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
# f3 Q- U. C0 r: V9 M; }exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

+ ~1 |# A1 a- v4 j# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

3 p# A1 P9 h, h5 @- I2 X0 _5 B$ e  T6 测试模型

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

def label2image(pred):
9 H7 B3 j  ?+ h% C3 g. M+ }    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
0 J7 [* O$ r/ C, L8 X0 h  K    return (colormap[x,:]).data.cpu().numpy()

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
9 `  J: [! ]$ p1 S) d& ?std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
def visualize_model(model:nn.Module, num_images=4):
: b7 G5 ~( ~/ I    was_training = model.training
    model.eval()
    images_so_far = 0
, b2 e4 R* ~, t: @0 P" q- [" w    n, imgs = num_images, []
  N( Y- X( D. z, |, t, C0 I( C9 a    with torch.no_grad():
! F. y1 P; q' j* O& g        for i, (inputs, labels) in enumerate(dataloaders['val']):
- V# i) S' C0 E- Y9 R% c# Q            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 # 记得要变回去哦
, ~: G/ J4 I  b" }- B% I2 P
            for j in range(num_images):
7 d& m* [  ~* `" ~2 Z+ k1 _                images_so_far += 1
. k7 X+ i& X  C2 _, y7 [( w                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
: J  p8 [# J2 J* V1 c                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
# d% B" T* S' i) `) a+ T                if images_so_far == num_images:
                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)
% q# x/ h% z, i  K! ^7 p9 N& \
# 开始验证
visualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
1 x3 `( G/ ~6 C# S+ z+ |def predict(img, model):
    tsf = transforms.Compose([
" R( E5 k6 N. Q; S! h            transforms.ToTensor(), # 好像会自动转换channel
- x3 a9 G' \7 V% @1 H! @            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)
. T/ Z& e6 T0 D. u! t    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
' m7 [( E" @4 K: ^* i) Z) O    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
  U) f! \" ~2 |6 i
def evaluate(model:nn.Module):
% q6 Q  z, R) t- E3 w: ?7 Z    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
: |' K+ e) r! H" o6 D    n, imgs = 4, []
/ B, `$ R" ^* l1 ?8 _. r+ Z: k    for i in range(n):
. |- l# c# Z+ Y5 a# L5 d4 r+ n        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
( c, ]' F/ T6 D/ s        imgs += [xi, pred, yi]
& g2 N* d- Y% ]8 b4 P* q7 z- B    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
( ?" E4 v' W# S+ t
. R1 t) }# _$ j: z* ?! ^/ ]# 开始测试
& v7 z4 i% \! t  L& F& Pevaluate(model_ft)

7 结语

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

8 ^2 }$ `. Q% k% ^Epoch 0/2
----------
train Loss: 1.7844 Acc: 0.5835
$ J7 j7 s3 G# g. [val Loss: 1.1669 Acc: 0.6456
/ p4 s/ ]  a3 q! E6 m4 B- J# s+ |- d
# P5 M) y* C) r( `9 `: dEpoch 1/2
, w8 W  p' J& i8 y$ A----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929
! y! f# \$ `" N4 B
Epoch 2/2
) o: }% u; S, A- B----------
. l, r5 Q1 U, Z# rtrain Loss: 0.9578 Acc: 0.6706
4 e% @0 U/ l: [% P4 ?) L7 sval Loss: 0.8088 Acc: 0.6948

Training complete in 6m 37s
% h# W; y( B# Y% _
9 Y6 u7 o( r6 W" j# \
9 c' R  w% L9 y, |* l$ t" Y6 M
$ v2 N8 x- x1 E! U7 N
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
% j; w6 v9 G; x. _) h- g
对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。


6 }" G& E2 \6 ~# `0 h0 W语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
! z; n. I1 T! a6 o; J! e. F5 @! h
. n. W2 f5 [& K- u% |" H) `$ H. fDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
$ p' F* H$ c, ~6 Q2 j7 j" uGCN 通过全局卷积网络改进语义分割[论文]
% y4 H/ v6 E" }6 Q' u0 q0 r: ^UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
. `3 |3 l* J) E2 jSegNet 用于图像分段的深度卷积编码器-解码器架构。
8 M, p) f8 O) i7 [  ]还有（DUC，HDC）、PSPNet等。

8 d3 b9 y  \: u0 W2 O1 |+ X1 N常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
; X5 I' ~0 k% _( h& ~8 _' U
" A8 e" n% m; ^) f5 r对于损失函数，除了交叉熵误差，也可以用这些：

Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
3 F- y4 q0 N# I7 k0 t* h& dCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。

- G/ d5 {) i1 }' c8 H

9 j; k8 ~0 q+ e% j9 b
; q: X. r  E& T- J
————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
0 U; `: Y. F$ h3 a