Pytorch实战语义分割（VOC2012）

杨利霞

Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
% N: A( C$ N  ^: ^- D. M语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
8 W/ c2 z/ d! a8 }& ~; n. J
" d/ \7 z- r( e0 R" t# B
9 D/ ~* s/ ~2 @$ S语义分割中图像有关狗、猫和背景的标签
文章目录

1 图像分割和实例分割
2 Pascal VOC2012语义分割数据集
2.1 导入模块
+ e- w9 j/ ^! z6 U" B1 G2.2 下载数据集
* _- O; {( ?- f  b& ^3 Y- c7 Y2.3 可视化数据
6 C, i3 X8 m: E+ L$ B- z2.4 预处理数据
" ?3 P9 m$ x! f* S( @7 N3 自定义数据集类
3.1 数据集类
3.2 读取数据集
( g7 Q9 p) g0 T$ D$ f4 构造模型
4.1 预训练模型
4.2 修改成FCN
4.3 初始化转置卷积层
5 训练模型
6 测试模型
6.1 通用型
, L4 K* q0 Y+ r! l* w6.2 不通用
; _* q( X! D  g" Y; d. ]7 结语
( J. V! r# m3 R$ K0 Y3 ^* @1 图像分割和实例分割
6 x- b  R6 m+ \/ ]+ G- v
计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
6 O$ x, g6 X" E" T( |6 C
图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
% t. e, I2 ?, f( C* W
3 Y" l2 L1 Y  u2 Pascal VOC2012语义分割数据集
0 T0 l4 B  C' U. O9 w; @
2.1 导入模块
import time
* Y0 @- ~& R$ `9 `3 Jimport copy
import torch
from torch import optim, nn
import torch.nn.functional as F
import torchvision
8 B6 F7 _/ V% w7 N% t/ x8 |from torchvision import transforms
, M, d' U& S) Z2 O8 z1 `from torchvision.models import resnet18
import numpy as np
from matplotlib import pyplot as plt
from PIL import Image
import sys
; p1 y9 ?7 i8 G. Zsys.path.append("..")
! a" M) r! i7 f; W# j8 g& s4 Zfrom IPython import display
& z4 Z$ J1 M* C9 ]" E* O, Bfrom tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

3 |6 ~4 T0 C% {  a5 ^, t2.2 下载数据集

语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：

3 U2 {0 D1 B' Q/ p4 {0 [
ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
' \' s& q4 ^/ B$ w# c2.3 可视化数据
. }$ I# H$ s/ Z/ @
定义read_voc_images函数将输入图像和标签读进内存。

def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
    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
/ p" w5 K% L( Q3 q. w    if max_num is not None:
        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
/ b# {; x  H% o" }& k, N, W8 X) H    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
3 O% o0 D+ G6 Q        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
3 C: Y1 O& @- t    return features, labels # PIL image 0-255

定义可视化数据集的函数show_images
/ b  e" P$ z9 l$ g& z3 ^/ u# k
) |: I/ ?0 {' A: ]( e: \/ g6 |# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
4 z0 G0 }' u- S& ?8 J1 Q    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

8 n; @' q# `" o3 n" rdef show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
& N. P& \/ f- w    figsize = (num_cols * scale, num_rows * scale)
7 P9 t4 ~$ N0 Z( o7 S% I    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
        for j in range(num_cols):
/ z" X! B, V0 Z% z" y) L" D            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()
8 j3 w+ p$ V& A3 [1 {) N# }    return axes

  m4 B& c% k. _+ M定义可视化数据集的函数show_images。
: ?, E# p: v& s7 V
% X) d" T1 d; d$ j7 e# 这个函数可以不需要
; }9 D. Q, t  R7 j! r; v* X6 Jdef set_figsize(figsize=(3.5, 2.5)):
: ~0 v( T* E! O+ W1 _& d% j    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
! ^4 d) q( P1 f  v, ~  n/ U5 c    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
' \0 U' b9 O; z+ g! h
9 k& C( J" Q9 {" K0 D% _def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
5 G! I* X' e" m% a' h    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
8 b4 N* m, y% e1 Y$ u    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)
( Z# Y$ D# n+ S4 a1 o% w) r            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
) d  D4 n! a, {  t画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
  _, ~0 S# N( q9 v: z/ C5 ~6 ~. [
7 J6 M! d& h/ F" c# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
& J7 e0 x; e: f3 X& P* R+ B4 ^; e" Sn = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
& S4 Y" L9 O; [$ D' N" {
" E$ o6 m  A8 w/ o

列出标签中每个RGB颜色的值及其标注的类别。
4 A) F* E; u! d6 Q; t2 N# 标签中每个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],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
( k% x* r) O1 E8 \' |                [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]]
# 标签其标注的类别
4 ?' t4 o! ^+ M, @VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
. I# P/ u, a% D1 e/ C2 y               'diningtable', 'dog', 'horse', 'motorbike', 'person',
; R( c- V6 e3 w( I, d! Q               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
2 \, S+ A7 m6 q5 f( Z$ ~+ [; |( N6 f. M
有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
8 n9 F/ n% `( H+ y( T( n4 afor i, colormap in enumerate(VOC_COLORMAP):
/ v2 E# ?2 X2 @# K0 |4 D' x0 r4 |    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
" S( w2 @4 S; k    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i

# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
7 [* {" M/ ?7 U6 N, Y* i* ^    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

' H5 E. S5 @" q$ E9 K& i4 h; l可以打印一下结果

y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据
0 T; E) B5 Z% h1 [$ W0 a& I
9 H5 g/ `, L" {. v# M; {% _在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。


def voc_rand_crop(feature, label, height, width):
    """
2 o" i7 e: u4 k6 X) X' H6 ]    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
$ V+ I# o* w) d, M    Get parameters for ``crop`` for a random crop.
- x$ G+ j; _9 Q- U/ P    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
2 Q  X# o# m7 t) @  Y    Returns:
& r" _- y, B7 [# X+ r8 z  F        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))
8 ]1 S" `( h4 m: y    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label
) c6 @- ^& a+ R& A$ Q5 z
# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
1 d6 r2 R8 Y( nfor _ in range(n):
7 [, D1 m3 }( M* i8 E    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
# v. ?) v" e0 }5 p5 C% K) M0 Oshow_images(imgs[::2] + imgs[1::2], 2, n);



. {; Y2 {; i0 D$ S" s1 ^

  }* N# |5 p  L1 b4 y$ p. T* y
3 自定义数据集类
' p, P' {0 b" Z- e/ a/ W
3.1 数据集类

torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
) r5 P* h. o! T6 i5 W7 Q6 W4 d- W2 U$ @
__len__ 实现 len(dataset) 返还数据集的尺寸。
: U5 I/ I( J7 `0 T' S__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
- {; B9 H7 N( k( ~# G, Z- ?
class VOCSegDataset(torch.utils.data.Dataset):
9 ?; f0 S3 Y7 H% F  F    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
2 x2 R/ J% _- y7 h& {        """
        crop_size: (h, w)
4 P7 g8 n7 H, R) o4 M& A6 {/ L6 M. k        """
        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
  o8 Q3 e& t+ C6 n# v; y# _# d        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
0 l: Y8 c: q$ e' u- ?6 k, V! l# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
6 B: \9 s; T% V% X8 I        self.features = self.filter(features) # PIL image
% x8 c1 w) A( j- F- B* B        self.labels = self.filter(labels)     # PIL image
        self.colormap2label = colormap2label
: N+ c1 v- J" i& C& W) e1 v        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])]
7 e2 Y+ J: S9 f- T
% Z- K$ ~5 z2 L0 b" Q8 ~! O    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
$ ]6 r1 E5 ^! \* _9 D                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
; m/ C3 U% Z* ?- S* U# c  m+ C
    def __len__(self):
: r! ]5 s& t! v, w7 N5 n4 g        return len(self.features)
4 e  R1 j( i6 v. w0 F& o/ E9 W# h/ X# Z3.2 读取数据集
1 C8 B1 U5 ]* h
% G( T4 a8 w5 p$ }通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
* J! ~0 w5 V8 P1 [  [+ G5 S% T
batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
- d6 C) h  S- U; Z( [crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
' Q! e" h8 c$ d/ s1 v
# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)

# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
- C0 a4 p# J$ U8 E# ~' }5 V) Qnum_workers = 0 if sys.platform.startswith('win32') else 4
( ]( G  v* n/ ktrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
" `& V- @: C; B. q' m7 G# \                             num_workers=num_workers)
6 Y$ _6 k! L  v
- \  W0 O1 M5 B# p, e' a1 l# 方便封装，把训练集和验证集保存在dict里
' z  d, h3 P* n0 cdataloaders = {'train':train_iter, 'val':test_iter}
4 Z7 [5 d. B0 i, l+ b! `dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

& T& w  W$ T1 W& G  T6 B' r4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
" i, l$ y, p5 ^* a9 i: O$ ~2 j8 g
num_classes = 21 # 21分类，1个背景，20个物体
9 A5 [* E" R3 X; A# g% Smodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
for param in model_ft.parameters():
    param.requires_grad = False
! }: g* S* [1 l; x$ j4.2 修改成FCN

8 b# |" G0 w/ R8 Q8 ~全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
) K/ ]( ^: M- r# ?- o* \最终输出的通道包含了该空间位置像素的类别预测。

对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
9 C" N! c, j* {/ M, R
! q2 Z8 l  J( d  c: c8 k  {可以先打印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 倍即可。
+ ]+ W. E0 K* K5 U+ y$ X
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小

# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
& ]7 ~& W: h7 y9 O& j% S( I2 s
' I* P4 G1 {8 l* Z, l( ]+ i# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
3 w' ]: W4 {) K9 Q  g6 o#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
#     break
, o# s0 H+ ^9 @% J
* l  `" N% D+ v
" a* o- h+ @& M) D  }) i. V& N# P4.3 初始化转置卷积层
. x7 Q9 i8 U. r1 g2 @7 X# }
; r# B6 x/ G& \" P在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
0 N% m2 ~# p: t! z- x; T! I! ]在坐标 (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函数构造的卷积核的转置卷积层来实现。

4 Y3 E% s3 u5 B! f
# 双线性插值的上采样，用来初始化转置卷积层的卷积核
3 s; V! t; B+ q. |def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
6 |  Q, V  A' `, @    if kernel_size%2 == 1:
  C  q4 i. R  o! ?' |. Q0 q        center = factor-1
    else:
& A, Z' Q; ]/ s7 j- u! ]        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
) d& a/ [5 {. [0 x    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
2 f& [' M1 T! l3 f    weight[range(in_channels), range(out_channels), :, :] = filt
" |6 [4 `$ C" X6 B" C! s+ U    weight = torch.Tensor(weight)
    weight.requires_grad = True
# I0 u3 c) u/ d4 l: q    return weight
$ y. h' B. x: b% Z; U1 H' c" C

在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
0 A+ h; q3 r4 r% N5 ?6 [+ E5 |
( ~& u3 v# b9 Z' s! fnn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)

$ X; x9 J; t* Y
4 _' O) q1 y: `* H: A1 F
7 W+ H8 k' N9 v' T, a9 T" A7 ]; ?5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
" ~; U1 I8 I( _9 @! m5 O2 x- D    since = time.time()
8 g1 Z$ Q9 @5 _& A    best_model_wts = copy.deepcopy(model.state_dict())
# ]1 f0 Q% o+ ?: R" j- A: _    best_acc = 0.0
0 C$ B# Q  S5 X) R  \& J    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
4 Q3 x! z9 X7 J7 R9 z' ?        for phase in ['train', 'val']:
3 r" S4 x. b  Z            if phase == 'train':
                scheduler.step()
                model.train()
8 d/ y- Y+ `$ W8 }            else:
                model.eval()
            runing_loss = 0.0
6 S* e7 _) Q1 Q: a; \- k8 M# H            runing_corrects = 0.0
8 J7 k$ v6 v" r/ N; }, ^3 z            # 迭代一个epoch
% c" X3 G% k1 j5 R+ D% v# A            for inputs, labels in dataloaders[phase]:
( H3 _4 ?4 v  z" @4 i                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
& r" q+ r. y4 Q) Z7 _                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
9 U' R7 f1 ^8 t7 l$ \# [                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
3 {% S0 l# r( u, X                    if phase=='train':
                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
$ `; P6 ~; ]0 q' u& }
            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
: z% c2 k. d5 p  H! @( E            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
2 A: c  m# t# o0 E+ j  W) N                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
# r$ ]& h: ]4 r% Z" N% I% q' C                best_model_wts = copy.deepcopy(model.state_dict())
8 X# r# P2 i, u; U+ a        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)
+ f5 {8 p1 G  d- @2 j! i6 l    return model

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

3 G9 h8 |0 ]( B0 J; D/ kepochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
, j+ @. c3 S1 f; |6 Z% R+ woptimizer = 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)
) W/ ^5 r' m) T1 s
# 开始训练
  D5 s/ ?1 B' z, h% s! K& ^model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
& B: ]' u* r4 [3 H. Z
6 测试模型

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

def label2image(pred):
    # pred: [320,480]
, ?0 P+ |( x8 r' X7 i8 t    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
6 |2 X$ o+ P' j; N    return (colormap[x,:]).data.cpu().numpy()
; O0 I. i" [# s  C" _

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
( [6 F& C0 f5 O" [std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
: y9 g: ^" S. }1 M" j8 i  I, E% o* L7 Mdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
    model.eval()
    images_so_far = 0
" J$ O1 b% p, j4 |/ ^' C8 j    n, imgs = num_images, []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
' {' `* C) _4 }$ c7 X  [            outputs = model(inputs)
) C, B. R: }6 ~0 G0 s            pred = torch.argmax(outputs, dim=1) # [b,320,480]
            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦

" F. i4 e5 q2 _5 H            for j in range(num_images):
' _: h$ s" T% Q$ p) |+ }8 K                images_so_far += 1
7 m% D' c9 s% M1 B/ [3 \                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
* }" Z! B, b- M$ h  @                if images_so_far == num_images:
                    model.train(mode=was_training)
, Q& D( V7 ^; @+ J                    # 我已经固定了每次只显示4张图了，大家可以自己修改
& |5 g/ b& H, u8 ^3 F  r                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
& y) C! k, e( v1 v) v. Q                    return model.train(mode=was_training)

- s  a/ a* L4 P1 |# 开始验证
4 v0 ^1 x* b' Ovisualize_model(model_ft)
% E  K6 W$ e  X: @# z: W6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
; F0 d, y0 |2 J6 l9 C, P            transforms.ToTensor(), # 好像会自动转换channel
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
  l0 _7 a4 b* N2 [9 E    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):
3 B* O. I( \1 V, K    model.eval()
  {9 L9 G8 F* i2 |. U    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
4 W0 `: m4 T1 b/ U) \% ?1 X    n, imgs = 4, []
" U  `/ B2 q- Y    for i in range(n):
" D" I; ?4 y! g1 X' }0 e        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
/ d: l: V" K3 [2 H9 K6 V        pred = label2image(predict(xi, model))
5 j" N' L. E1 J$ H' Q        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
6 y7 i6 s( j" z3 A
7 C$ c4 l  T* V+ f( j4 k# 开始测试
: l- W% l' f% `evaluate(model_ft)
5 U5 c' }% ~4 }. u: t  B
/ u: T5 S, N& o( y( W7 结语

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

. A9 Y; j" Q/ ?% E* xEpoch 0/2
----------
7 r; S3 h8 q5 e4 m9 Wtrain Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456

! p( a3 c" G: R& T" kEpoch 1/2
----------
train Loss: 1.1288 Acc: 0.6535
3 P4 v. e7 @4 {7 S9 E2 rval Loss: 0.9012 Acc: 0.6929
( z. z( V% X, V% V' Z* D
Epoch 2/2
. ~6 X! G& {' B----------
1 o( ~; m3 R$ Xtrain Loss: 0.9578 Acc: 0.6706
: h* x9 [( h; p# s! G, a- i) y' U( ival Loss: 0.8088 Acc: 0.6948

: t0 T% Q- [( N% g$ k3 T) `, RTraining complete in 6m 37s
: z6 n. C$ I, w

1 t' k( E. e. a

. r0 k7 Y$ m8 K3 J$ e当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
& m( r# A. W( Z* P6 v
对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。


5 q# C% j3 K* a语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
8 B, E  u; p% SENet 用于实时语义分割的深度神经网络体系结构[论文]
6 [  Z0 r0 ~6 M$ T; h2 [" [U-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。

常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
5 Q$ T4 d" X( S3 B; k; ^9 B5 t6 q
1 h( S" J0 W# T- g* o! N对于损失函数，除了交叉熵误差，也可以用这些：
- Q; O8 K4 t! g- r) O0 c
Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
& l* G. T7 S9 {Lovasz Softmax 查看论文：Lovasz - softmax损失。


; O: a/ W* y; B+ U& ]. p% V
+ t' V, z! ^- z; m# Z

) ?' d. I) X" n————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
% {1 g' J# H- x+ S5 |) ]原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507

$ Y' o5 |: G  _6 f# N
. G9 ]' Z2 H& ?8 \