Pytorch实战语义分割（VOC2012）

杨利霞

, ^1 |* J6 I! j% x( Z/ sPytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。

3 m) T( T7 n& {* T- P
语义分割中图像有关狗、猫和背景的标签
) g: d5 d& ]4 m8 T% a文章目录
  o7 j# L: ~& S: M
1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
; w3 F. _; Y1 _% E; T6 U6 b2.1 导入模块
5 A3 c: m, h) j( S7 b8 G8 r2.2 下载数据集
& _# }. X$ p8 ]7 G) s" a2 ?2.3 可视化数据
4 b- u& a5 I. b. W. R( \; U8 x2.4 预处理数据
- ?" k3 T; W6 L. K3 O3 自定义数据集类
3.1 数据集类
+ ]" ~( g6 `; r% Y/ e0 f( f3.2 读取数据集
' T! V; g1 {! P. }- z4 构造模型
8 |3 t9 u: ~# i# b2 d4.1 预训练模型
4.2 修改成FCN
; H, l' ?: p, c. z  e4.3 初始化转置卷积层
5 训练模型
6 测试模型
0 S  u2 j( I) I: v% r- `+ d& ~6.1 通用型
* H, u: n5 O0 ^$ Y3 Q6.2 不通用
7 结语
1 图像分割和实例分割

计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
4 L9 z1 m' x5 l! v
" p+ z8 P( z; r$ v( K! i图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
+ f& t5 a" g+ g5 y0 B
" I- u' Y$ }  O+ N' i6 F" r6 y, j2 Pascal VOC2012语义分割数据集
, V# W  N% ^7 ]* G$ n, W/ o/ E2 v$ a
2.1 导入模块
import time
! H) L7 e# W# R+ Dimport copy
import torch
from torch import optim, nn
5 q2 f, x/ @. V( Pimport torch.nn.functional as F
6 w( b. B% z% s. himport torchvision
from torchvision import transforms
from torchvision.models import resnet18
8 i4 C/ |* K3 Q4 G9 F' C5 }) |import numpy as np
from matplotlib import pyplot as plt
from PIL import Image
import sys
sys.path.append("..")
* W" Q/ S2 {  Vfrom IPython import display
from tqdm import tqdm
import warnings
- y9 p( r, w% d) V% M# [) x( Q! M; Y" bwarnings.filterwarnings("ignore")
, g' @$ I& S  k2 B% ?
2.2 下载数据集
4 D7 I% i# N  t$ s
语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：

- [# \, q8 G- s( {3 X! j
3 P/ b# I! E( c3 S7 q+ c6 C4 a" i' gImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
$ v: \% x: J8 A: f- ?JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据
! `' N0 I4 I( d# j$ K& F
) V" M! s* `- E  S& @+ J: l) i定义read_voc_images函数将输入图像和标签读进内存。
4 b* R+ b" A. M. ~% I+ H/ E
9 ~3 |3 w5 ~3 q# xdef read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
* Y$ m& ]5 N/ [4 `    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
% O# W- u$ j0 H2 R* @* X    with open(txt_fname, 'r') as f:
( n& k# Z7 b: ]- r( p3 K        images = f.read().split() # 拆分成一个个名字组成list
' A1 ]# K/ o& b5 m, ^1 j6 L* j    if max_num is not None:
        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
9 H! I9 l+ k! x& ?4 B8 t' m        # 读入数据并且转为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")
3 B3 p+ E7 c# a! ^1 W0 H, H8 L    return features, labels # PIL image 0-255
. _4 ]9 j* d# ]1 _: l
4 B) z; x/ G5 J9 X5 @  V  E6 T0 ^定义可视化数据集的函数show_images

7 E- R1 j; c& u% F( \# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    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)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
, K% U& q' }" D. g: X0 N2 G        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
            axes[j].axes.get_xaxis().set_visible(False)
! g2 F' {# J0 B! O            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
- o6 H4 \. e4 K6 B# y. a0 J    return axes
! y5 h0 a9 d, `* H9 ^
* Z2 F, P! o1 U9 O: q( g定义可视化数据集的函数show_images。
' |1 q% g- o8 ]; A% E
# 这个函数可以不需要
' o. V$ P' y. J" ~def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
% p$ y' N* Z" m4 _4 F1 ]    plt.rcParams['figure.figsize'] = figsize
$ F& j9 H0 x! d
def show_images(imgs, num_rows, num_cols, scale=2):
9 ?: z6 T1 H% X5 D$ B5 L    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
1 l8 e6 s; p& X" o2 G        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
$ ]( P- M: R! n* U  E            axes[j].axes.get_xaxis().set_visible(False)
            axes[j].axes.get_yaxis().set_visible(False)
& A! @2 ?# k' P6 R" Z1 ]    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
: h& u, M8 C# H- q: e3 L5 f
# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
  {/ N- C6 M" S  c2 Pshow_images(imgs, 2, n)
) f& _, J: Z/ `, s, S) F+ k
+ `! X, Q) n+ x3 x& W7 t2 f

) ]/ A. I* |7 K3 E6 D9 Q列出标签中每个RGB颜色的值及其标注的类别。
+ S1 Y" z+ J* z  C0 T, J8 y" u# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
3 @* E# V. a% ]                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
- o. v. ^0 w' W6 H  B                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
- L; Y/ s1 I" M& q; T                [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]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
2 x( z, X9 F4 Y! W0 f9 d  v" I               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
' ]  B* \* a# w9 x) T2 k. [
8 {& C3 k- {$ O: e' f! x有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
" t( C# e( \9 o9 b: `colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
for i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
# P, a4 d7 e! D6 K! i2 ]( t* ~    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

9 A# c; y4 o4 V& n# I! [# 构造标签矩阵
' J5 S- H: F2 hdef voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
7 ~+ U& K5 K" g    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
) ]" d  X* s% R$ T- \    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
6 I) @8 r1 N! }7 @) L
可以打印一下结果
9 j2 ]' V0 f& V' h
: p7 }( }% y0 hy = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
2 Q; ^6 `* G" K  T
9 F+ w+ S8 C* P( h- D2.4 预处理数据
& y1 J0 ?' K( \' E0 R
在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
. L+ P! s8 n% F0 c' Z9 y( X

4 ^; @5 {0 {6 h! Y, s& T2 w# K4 Odef voc_rand_crop(feature, label, height, width):
+ ?/ h" [( G) w" Q6 Y    """
    随机裁剪feature(PIL image) 和 label(PIL image).
# A# W% l6 p3 {    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
. V4 [1 g" N& Z$ N* x    Returns:
" C) ^! f, l- G5 X        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)
; s! o* D- w0 E6 |    return feature, label
' W0 S% u# [9 s/ h8 m
0 w3 d* F! N% J, e2 d# 显示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);

1 n0 Z' C  E& i+ S) d
  t5 P5 z( ^; o- n8 p$ w7 G

. g& {7 p2 p2 w% M1 ]3 x* J* I

3 自定义数据集类
. T' y7 ~8 ~* V
3.1 数据集类
  |, f, ?. M: l: k' d
) m$ W! y+ q* d) \3 ytorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
% s1 A/ g% x3 t
0 r/ {) M7 s: c5 u4 m__len__ 实现 len(dataset) 返还数据集的尺寸。
% s5 R4 C3 X3 N0 [6 N3 Y8 b__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
& c+ j' W# S" l! ~9 X由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

, Y0 c! {5 T4 b2 m6 B$ j  X, eclass VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
: C% m' F3 A. @+ K4 g        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([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
' Z6 b3 ^. d" O! V4 W9 q+ o  G        self.crop_size = crop_size # (h, w)
  t8 N3 P& S3 {$ a        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
6 {! \, a9 S! v# k; P; b9 t2 |# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
5 o: `  z0 P7 Z        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
- g- ?! u5 O- s- G+ m$ w- n4 p        self.colormap2label = colormap2label
( \: \% k3 v1 i! W/ I8 w6 C5 E        print('read ' + str(len(self.features)) + ' valid examples')

' ]8 A2 @4 C) @    def filter(self, imgs):
& B) x& `* y$ Z- N        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
! }7 K" w5 |( Q2 Q* k; f( _: |0 b) h/ t
) i/ ^+ v) O$ W; z+ @9 s% I+ i' r    def __getitem__(self, idx):
  H. s* |, R3 c        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
                                # float32 tensor           uint8 tensor (b,h,w)
; z  A4 N$ ^" F$ O6 Y- ~, A0 j        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

" H' }; W" B6 @& j: D    def __len__(self):
0 x- R- z3 m' G4 Z, h9 W7 `        return len(self.features)
3.2 读取数据集
, I/ R! V) N* ~6 ?- E6 {
通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

) K" m; Y9 m7 p: J) L, abatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
+ V! G; D5 M. W, W- h: ncrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

# 创建训练集和测试集的实例
5 a+ C$ O, q* k+ e! ~! Ovoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
1 }  q8 `+ d. m$ f$ y
( O/ m* o$ w* C+ {; `) O# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
& ?0 T2 f! Y" O  atest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
: ]0 Y" @# d/ L2 w. @. T                             num_workers=num_workers)

- ^% V6 E( n% k# 方便封装，把训练集和验证集保存在dict里
( Y# r% `! _5 j  Ldataloaders = {'train':train_iter, 'val':test_iter}
- L) D$ }- i4 M/ n: R3 \* Z( Vdataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}
: T: ^( }/ Y4 z! I+ ^% ]1 t
$ W: A, g9 o! R3 G2 }, Y4 构造模型4.1 预训练模型

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

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

num_classes = 21 # 21分类，1个背景，20个物体
9 p! J" W% Q4 C' f  T; \model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
for param in model_ft.parameters():
    param.requires_grad = False
8 b( v+ y+ B+ E# y# t: \- B% }- ]4.2 修改成FCN

/ i1 E. t( L0 ^5 z" Z( g# K+ w全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
最终输出的通道包含了该空间位置像素的类别预测。

对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

! k6 C5 T; `9 A# A! K$ g5 u/ G" ?可以先打印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 倍即可。

model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
: D& g2 r; J4 h( T5 N& D) d              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) # 转置卷积层使图像变为输入图像的大小

* ~# L- }, Q  C8 @9 y# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

+ h( P! r: r6 r# z+ C: i# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
, \* l- W$ H& d( j# for X, Y in train_iter:
! e9 y. V# E- [/ ^+ Z  K3 D#     print(X.dtype, X.shape)
' H/ z9 j% r* N. _7 ]#     print(Y.dtype, Y.shape)
#     break
2 ]  t- l3 F1 ^$ x" C

4 @/ J5 l6 Z9 U( e4.3 初始化转置卷积层

在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
5 E; x  L+ o# f$ M/ P! x) Q, @7 C. f在坐标 (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函数构造的卷积核的转置卷积层来实现。
6 k* F( S/ C0 l) J4 U" A! `
$ G9 a8 X# A7 p7 ^
* M  C- l2 J4 i* x  Z; M9 ^, U# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
) \( Z/ m# l# ]8 d  K4 q' h    if kernel_size%2 == 1:
5 I) W1 N& @: {/ J7 G6 M( I        center = factor-1
    else:
6 V. D6 Q+ g- w( ?9 Z# N: W        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
5 T; a; Q# G, j3 n    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')
: h& r& ~1 l; g& E: l    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight
0 n5 U9 b) c- Z' M( j# c
9 m4 P& E6 E; y/ k8 M( J# f9 U$ M
在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
5 O' x9 r( S# l8 Z1 i
nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)
9 e, n3 c* @! R* w; E( z

1 C0 o9 {5 [. T
5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
: l7 T3 ~! @1 ?2 w" v: U    since = time.time()
8 I9 o) R% G* q. @, \( J" v    best_model_wts = copy.deepcopy(model.state_dict())
, N$ V: P6 L8 \* z( m4 m* P) m    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
- D# {8 B, x) i  k6 u5 U+ b        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
        for phase in ['train', 'val']:
            if phase == 'train':
+ x5 S  s  }- |7 A                scheduler.step()
( M5 d( i1 `' t+ T                model.train()
            else:
                model.eval()
            runing_loss = 0.0
. T. r4 F* _, y& P# g8 _; F2 x" K, L            runing_corrects = 0.0
            # 迭代一个epoch
. `3 C0 L* p" |- ^            for inputs, labels in dataloaders[phase]:
" P# H( u, o! G* x2 r                inputs, labels = inputs.to(device), labels.to(device)
( Z$ j5 d% k; T6 _: h                optimizer.zero_grad() # 零参数梯度
8 g3 f7 u" i: _* Q- z( P" T                                # 前向，只在训练时跟踪参数
5 g) N8 R8 c, ]- p                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
6 N6 O9 }1 r6 ?/ D/ M. @                    loss = criteon(logits, labels.long())
  s  z2 l# ^  N7 u                    # 后向，只在训练阶段进行优化
# r) I+ d' @! z/ A8 D                    if phase=='train':
                        loss.backward()
- E" |* M( m7 ?5 u) m                        optimizer.step()
4 K; m$ `" T' V2 i3 q0 A8 u4 j9 {                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
0 s$ k% ?& }$ Y% M            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
- b9 N% E& Z: m; e: |                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
/ [9 W( H( m4 ~8 c. X" ]                best_acc = epoch_acc
( D* x$ i5 z6 b$ ~( l. ]3 E                best_model_wts = copy.deepcopy(model.state_dict())
, e+ ]. Q3 M+ G7 [        print()
    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model

下面定义train_model要用到的参数，开始训练

epochs = 5 # 训练5个epoch
# Y* ?' Q' A1 r7 f" ~4 Fcriteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
) X% u0 }' q5 S# G+ r% o2 |3 S# 每3个epochs衰减LR通过设置gamma=0.1
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
4 @2 D5 I; Z$ M! N" e% X7 `
# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
# B- h7 G- Q8 G% ^( d6 o
2 S/ C( n5 E9 U3 N* w0 D6 测试模型

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

def label2image(pred):
    # pred: [320,480]
/ n& S3 B# E) q- p    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()
+ v! ~4 g+ w+ e1 `4 V  p* w/ ^

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
0 \0 _' |& B( H( j2 O' O, X% ~9 }! U6 ]def visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
    model.eval()
( d6 U/ H% l! x5 P6 b    images_so_far = 0
    n, imgs = num_images, []
    with torch.no_grad():
8 i! }' ?- U! m# e$ J        for i, (inputs, labels) in enumerate(dataloaders['val']):
            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 # 记得要变回去哦
1 |0 q! P; V$ m5 m5 [/ y6 s
            for j in range(num_images):
: Z3 c' O0 i1 |; |                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)
$ l6 p/ f) q! I                    # 我已经固定了每次只显示4张图了，大家可以自己修改
! F$ N$ y9 S5 q) ]# y1 x, K                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)
$ U2 M( e8 X  c. f7 Y
# 开始验证
visualize_model(model_ft)
) ^  |, A$ I4 f! O* c' ]" O6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
9 H$ m6 [8 v% [* ]( ldef predict(img, model):
    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
9 A$ b0 j' {( ?3 e% c& R            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)
! w# l3 p+ b8 I1 o    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
( ~9 g4 Z; I4 Q    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
& H$ n7 T# u- Y, r% y3 f! R& P: w. Y
def evaluate(model:nn.Module):
    model.eval()
    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
) {! Y" E3 z9 D' \( `$ Z/ t- X# g    n, imgs = 4, []
    for i in range(n):
. S# V. o  e0 _+ h) L        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
7 U4 j8 z+ }, o        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
# i  M# `% l  E. l- e9 I    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
; h% [' x; M: g% X
6 }3 {3 ?5 g! u# 开始测试
evaluate(model_ft)
6 R3 N& y& W, j0 S9 k! n. G. o
7 结语

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

) v/ i) B1 ]; U: l. l7 h4 CEpoch 0/2
----------
0 q# y4 s  R) _& a# |& Atrain Loss: 1.7844 Acc: 0.5835
2 v* D  U; B  @val Loss: 1.1669 Acc: 0.6456

0 h  y+ n6 L, \3 d; x; bEpoch 1/2
7 d. Q& h2 L$ V% i, \3 X----------
train Loss: 1.1288 Acc: 0.6535
0 @# y* \7 ]: O7 }/ k. i& |1 Z9 jval Loss: 0.9012 Acc: 0.6929
& x  l& |: I) O8 U; A9 N
Epoch 2/2
& a0 D( u( u, y  P  e3 R$ Z----------
0 }0 T! {! D! r- X; E9 Vtrain Loss: 0.9578 Acc: 0.6706
8 v( H4 T1 q; s" vval Loss: 0.8088 Acc: 0.6948

Training complete in 6m 37s
9 }; D9 |) {/ R' B) L" T' ?  w) c
2 S9 K1 A: U6 P( @; J5 t


4 \+ N0 h1 T5 t; m$ r当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
# Y) h: l, h3 p
% ^! c% {# \# h0 U$ ~对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
2 Z$ q! s5 g5 K+ x* X' e

, T' V; i+ ?5 f/ x& B. |/ x# y语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
4 y: N7 u2 p3 m2 G- f8 j( c
9 f! k8 ~5 v) ~) A* fDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
+ |8 j* N* X! q5 h& B! r) BGCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
! i% `; A8 I8 T& [& r% ^ENet 用于实时语义分割的深度神经网络体系结构[论文]
6 m- e% W5 X* I2 w/ rU-Net 用于生物医学图像分割的卷积网络
4 r, u; F4 T  O. HSegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

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

对于损失函数，除了交叉熵误差，也可以用这些：
- u$ v$ g3 k: a  e
8 U' }* @7 T* K: ~9 {3 a  ^Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
1 U" A) R! a% j: f5 [5 {) D4 l2 eCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
' }" b2 p( Y; G  h# KLovasz Softmax 查看论文：Lovasz - softmax损失。


0 g$ R+ b2 }$ }) @8 X4 K
0 N7 `3 h' j3 s8 l% h  g

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

$ b; V; y. U. T) H3 y