Pytorch实战语义分割（VOC2012）

杨利霞

; Y7 X5 g' a# p- t
Pytorch实战语义分割（VOC2012）
! ?1 C8 o' m' `9 x本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
( [& a/ U( {$ [, f语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
* m9 E% w) y5 H  F: h9 r+ ~
0 t' S; l' A9 C5 j6 K
语义分割中图像有关狗、猫和背景的标签
文章目录
! `+ J# h  [% P) u9 u, \5 n2 D; ~
6 y4 p% V$ F) _; J1 O* [3 ^7 Y1 图像分割和实例分割
/ y( y- ]) z) F/ e2 Pascal VOC2012语义分割数据集
2.1 导入模块
- n5 \. q, [  [5 T) }2.2 下载数据集
2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
0 t4 ?( v$ L( ^3.1 数据集类
# f5 i) x. N( ]! U2 m3.2 读取数据集
4 构造模型
/ v5 g9 }. a9 ]5 a  E# J+ E% G4.1 预训练模型
* f8 }: r: a2 {0 `/ B4.2 修改成FCN
( W/ E. F* Z9 r4 k+ i% H/ k4.3 初始化转置卷积层
9 }# T& _% m/ ]3 J' r8 c/ F5 训练模型
1 c. g6 V/ t; k/ ~$ i6 测试模型
# v0 ]& ~/ g, o) m4 f5 t5 P6.1 通用型
/ w( c3 f; A2 {9 O" |: Q# c6.2 不通用
7 结语
1 图像分割和实例分割
6 ?7 P+ Q4 b* G1 d) z% x
计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
3 l1 o' }9 D% G- X
0 ]# U% K, [: k图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
  a- |0 @( n3 B2 O
2 Pascal VOC2012语义分割数据集

2.1 导入模块
import time
import copy
import torch
from torch import optim, nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision.models import resnet18
import numpy as np
from matplotlib import pyplot as plt
& x! ^5 q6 S3 y* Zfrom PIL import Image
import sys
sys.path.append("..")
  {! |- t0 ?. n, k0 ?from IPython import display
% Z! ]4 s7 v  N1 W( z% }  ^% Ofrom tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

2.2 下载数据集

语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
# g, D; j, i3 X
% Y* {% @5 T4 }- t
ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
9 J3 s& K) w, a9 qJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据
) M6 ]* [, d, ^% D7 V
3 ^  h( u' P* I3 F$ d& g# a定义read_voc_images函数将输入图像和标签读进内存。

6 d3 w5 p% t0 Q; [. {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:
5 B" @$ l: o7 R6 D% j        images = f.read().split() # 拆分成一个个名字组成list
    if max_num is not None:
, U5 b- e' x6 `9 ?; C9 ?& P1 X4 \/ V        images = images[:min(max_num, len(images))]
    features, labels = [None] * len(images), [None] * len(images)
* G6 @: v  M# O, S( [' j    for i, fname in tqdm(enumerate(images)):
        # 读入数据并且转为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
: ^5 D7 W' v5 {6 _4 I6 |
定义可视化数据集的函数show_images
- T$ q% d; C& v# ^
# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
; e, q. ^8 G$ F1 v+ K8 b9 _8 t    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
+ D/ r- d' h1 p0 v6 k) o% y1 j
def show_images(imgs, num_rows, num_cols, scale=2):
' j2 ?, J' {% ~    # a_img = np.asarray(imgs)
7 g) `/ S  ^2 N. C    figsize = (num_cols * scale, num_rows * scale)
: |2 ]3 O' a# [5 d( F; l2 y& w9 p    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
7 j0 x! [) }. b; o4 [    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)
            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
* o5 _6 Q4 n. a% R# c: j    return axes
0 b5 h) ^. q6 q' i3 R
, m6 n* D$ c: D; [# @. T" k定义可视化数据集的函数show_images。

1 v# h% G4 j, a0 v# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
8 o$ K7 W1 B! G3 Q    """在jupyter使用svg显示"""
; \: M6 e4 e$ }' P5 m    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)
! d# [: c' b0 I, u! S    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
9 E& F! ?. s, d1 Q( r/ [    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)
* u2 s5 E) ]! e' p' J- O' {0 V6 t            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
6 D: f+ ~4 o. q/ H$ Y8 T    return axes
/ E) h# t; w( M) O4 o; q$ j  T画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# 根据自己存放数据集的路径修改voc_dir
- ?/ d& F/ u, Ovoc_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
. S" f$ b$ l2 {9 \show_images(imgs, 2, n)
. Y: X2 ]7 w: B3 X' b- x8 |
8 C* K/ Q1 O( \" n

3 x$ A* L" d! A' a' b列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
. |) U* f' E+ FVOC_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],
/ X1 R2 |: U9 Q* v) t( p                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
( V% f6 _# ]7 g+ Y, T/ W" t                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
! m" A. o6 P1 z* w  x5 `+ A4 l4 C                [0, 64, 128]]
1 ]/ S! A2 z8 [3 o5 v5 O8 }# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
+ {0 H3 f$ o: Q/ S/ k- q               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
9 L; e- N: K+ K" l# j/ E7 O# g有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

2 {2 g# V" q9 J0 ~" V有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
" P. Y" |1 k" L' e* a  bcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
$ N, M: A- [0 l8 j: L& a7 vfor i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
4 |1 Y% D; ?5 L) {    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
* X# q* u  S% }+ v7 G0 o# S
# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
! `$ R' a' I6 N4 C. E% W2 U
可以打印一下结果
7 ~4 P. B* h( X; Y9 b. \9 C0 y
/ T( Y- r: g2 {8 L& m/ W$ My = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

3 k( I& M$ F& a2.4 预处理数据
2 @8 C. o3 T$ m6 [: P5 T& s) T5 }
7 t* T9 R4 i- A8 m在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
) c& l3 \1 X$ F5 ]% {
1 X" C: E6 C  [; g' P8 g
/ _1 C+ F: u" \def voc_rand_crop(feature, label, height, width):
    """
- G5 H. k% s0 {0 i    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
    Get parameters for ``crop`` for a random crop.
    Args:
4 l2 Z7 {; V0 k+ C, L        img (PIL Image): Image to be cropped.
7 v4 h2 Q1 W! G0 H7 ?        output_size (tuple): Expected output size of the crop.
% ^1 V+ R0 e0 |" H' B( G    Returns:
' Y0 `( v$ j: D# U! I        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
    """
0 _5 `9 u+ c, F8 ?% x; f' }  `; R    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
0 u- w' u/ U: A2 ~) E    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label
: T! ~4 d( u$ w
# 显示n张随机裁剪的图像和标签，前面的n是5
9 P+ j7 L, N; [- N( Limgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
+ Q" u$ j0 \- q; k" ?; o! d( Lshow_images(imgs[::2] + imgs[1::2], 2, n);



) i. P1 o) ^0 P, c& X


( Q0 \7 A+ U# h# p- ?3 自定义数据集类
8 N% P( c# i1 x. m; |" h& B
# F$ h# _& n. e" o9 R7 E$ \3.1 数据集类

) j2 w: u1 E! m# itorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

- g8 ~" M4 w% F+ p7 p__len__ 实现 len(dataset) 返还数据集的尺寸。
" |9 g6 n9 @2 W; m& B% y__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

class VOCSegDataset(torch.utils.data.Dataset):
2 f+ m' ~3 h  M. t, O    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
        """
  s7 F: Y& a$ e        # 对输入图像的RGB三个通道的值分别做标准化
        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(),
  z, X0 f1 _% p1 ~0 L/ J3 R; h            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
6 {  X" p$ Q/ s1 ]3 \# a        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
. E- N( s2 F1 e$ F7 q4 K% f        self.labels = self.filter(labels)     # PIL image
8 D+ J' E) x7 X2 g) J6 o; T# P8 @, I        self.colormap2label = colormap2label
& R7 M# l3 o  v' K5 N        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):
8 }$ t* P/ P" M        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
9 x9 a: J4 y4 f7 t* K                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
- l$ n: P2 E* j$ m
3 F  Y4 z5 {- g  K    def __len__(self):
0 [% F) P' ]. N        return len(self.features)
3.2 读取数据集

0 o) i- z9 }1 `2 Y, d通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
1 x3 e* ]8 I; Q2 b3 C
% L8 r5 |2 q; T3 Sbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
& v8 [2 v* a# K3 P! n2 Wmax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
2 p. }3 g) d) |9 X  `3 Fvoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
7 A, ^7 Z1 g$ L8 @
6 F0 P# ^# j) C# 设批量大小为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)
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)}
, o6 V$ y0 S/ U$ I* I
, J2 R. S+ p: b0 y* u# S& ]4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
: u2 b( }3 c& e' d0 N
2 h7 i" j+ J! A" wnum_classes = 21 # 21分类，1个背景，20个物体
; W3 v; m0 p6 ?4 T8 D# j, ^model_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
2 a7 d& ?4 q2 f3 |; G; O
# 特征提取器
for param in model_ft.parameters():
, N* X3 U( H% R5 X4 D7 a3 H    param.requires_grad = False
4.2 修改成FCN

. v5 r- g1 R0 E- S2 e) U# c% W全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
4 i. r4 y4 G" k最终输出的通道包含了该空间位置像素的类别预测。
2 C" k$ w: N, e6 M4 J
* N) l% Z( @. K' k对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。

可以先打印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 倍即可。

% ]+ Y  _* t8 L& V2 H4 kmodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
# {" G& ?! E1 x4 K4 S              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
$ l! T8 v2 d5 \              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
4 A; g1 n9 I* m4 J$ w
- S# y5 O- P2 x! [, F  z# 对model_ft做一个测试
; D% h; n* k4 l0 bx = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
6 H! L; B  S: M5 _* [print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
% F7 l& H% N/ n  N! F
1 g. }+ Z8 W$ e6 e# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
) f+ l; X1 M; F( r; u- W#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
#     break


4.3 初始化转置卷积层

: \9 @! e: D2 [$ G5 O% ^5 s在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
1 e0 ^+ F2 Y5 [$ F' a9 i! n0 ~  Q在坐标 (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 Q9 Z5 x6 U0 @
' C7 b  v: G2 V* u8 w  \
, ?7 a: ?! \% Y6 M2 t3 v8 Y& C# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
        center = factor-1
1 r  d. `2 c+ u* \7 A( n" @+ r( |    else:
; p3 G) C+ s4 i$ }  H( Y/ r) |        center = factor-0.5
; J- E4 e1 |, T+ q    og = np.ogrid[:kernel_size, :kernel_size]
, Y5 d0 V9 [; {: t+ C    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')
& i8 i1 Q2 i% c+ M4 ]& K" U6 `- q    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
    return weight

; N# r$ a! c; K# \& ^' F7 W, M
在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

# ?0 j" u) j' D+ R2 O  s- J2 unn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)

( D+ d5 `7 E7 G% O0 [

- ^4 g5 ?9 S* O' w. b5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
; `+ T2 N5 l% p" F9 j    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
  w/ z$ x, E3 t, l8 b    # 每个epoch都有一个训练和验证阶段
! s4 O- i4 _; y% \    for epoch in range(num_epochs):
0 w+ K: j* E& a) q        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
: A0 S. {% f( ^7 S+ r! \* u2 N( u& E        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
+ O0 v* E% E9 r7 k, c1 f! S                model.train()
1 w7 O1 s1 @7 N/ |) t& o- ~            else:
                model.eval()
            runing_loss = 0.0
7 J. `( J/ f8 X! }0 i2 K            runing_corrects = 0.0
& w" W+ X( A3 U$ G% G6 v- Z            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
( s/ c' e1 [( Y2 ~                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
3 D" C, O9 G% h: R                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
3 i. c4 J' z5 ?, ?: Q                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
                    if phase=='train':
5 @3 v! E7 J+ U  B: ^                        loss.backward()
                        optimizer.step()
4 K+ h1 z9 w; o                                # 统计loss和correct
- x4 V$ X2 g( o+ l                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
: Q; n2 T! A1 _) O6 Y! L5 c- b
            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
+ L* p3 ^3 i: d- `/ t0 f            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
" `# I$ W- j" _( a9 f+ G1 W( H& m                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
5 I7 e4 N& ?0 Z( J2 P                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
& m( g! c" u4 ]0 l8 `        print()
    time_elapsed = time.time() - since;
6 T; Y$ @7 K8 v# N8 y' }5 A    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
0 I8 y7 J- v5 V* D, ^- Q    # 加载最佳模型权重
9 P# w3 Z/ f6 G( _    model.load_state_dict(best_model_wts)
1 M; T- s( o9 p  g    return model

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

epochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
+ \; u- J( F$ w. v6 @9 A# 每3个epochs衰减LR通过设置gamma=0.1
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

6 ^6 ]: W# T; N. D0 P# 开始训练
0 [7 ^4 `1 q  u1 _8 y$ Rmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
9 e$ x; d& z; R. H# g
6 测试模型

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

def label2image(pred):
3 g2 O3 ?# x; n! D    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
& B( k+ s% R% T    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()
, Z% u7 n* d: R: C* k$ Q: q

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
& P% F! Y- P: w% Z: t8 `std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
$ z& g' l9 O) N, Y4 n+ qdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
& b, y. I4 J% c+ K3 I% r& P! F: c  G    model.eval()
/ Q3 I' l. K& h/ @% M# U    images_so_far = 0
! ^) ^+ t  x3 w  v; D& M    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]
            outputs = model(inputs)
+ b' d  b9 c/ G) _) D: D0 z            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
: x' p4 P7 t: w) d% a% ?                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
8 ?- I  h6 A$ H) N5 S                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
, c1 z4 x8 k0 E# q4 x                if images_so_far == num_images:
                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
, w4 g0 t6 e7 ?                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
) P, r. J5 |! Q                    return model.train(mode=was_training)
5 o2 ]- ?1 N7 U% X+ D
# 开始验证
visualize_model(model_ft)
- u, C1 e2 |( Y, u: G1 @) M+ i6 B6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
0 h- J3 e+ A8 F& a; Y3 J            transforms.ToTensor(), # 好像会自动转换channel
            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)
" x1 D& J4 [: g2 \% ]! b    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
9 q0 G4 M) G3 u1 u3 }# `9 Y
def evaluate(model:nn.Module):
    model.eval()
# k2 a" I7 X) u    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
    n, imgs = 4, []
    for i in range(n):
        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
' {9 M" E+ Q, Y2 N) Y* \+ Q0 ^# }        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

# 开始测试
evaluate(model_ft)

7 结语

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

5 @, w8 O  B3 s2 Y1 w0 K1 y+ B/ w; BEpoch 0/2
----------
train Loss: 1.7844 Acc: 0.5835
! o+ i5 o! m5 I2 w, _val Loss: 1.1669 Acc: 0.6456

: \! |. i  E' p* e4 H+ kEpoch 1/2
: ]  A6 N6 h: l! P----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

Epoch 2/2
0 y3 \/ @; w( L9 L5 F+ `----------
! J0 R' ]( d  R# w2 Z3 s6 Xtrain Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948
1 x* p  R& V* O9 x; i
Training complete in 6m 37s



! g2 d- s( U1 X) U
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。
, _( R- k* C( x( ^- g- C' u
对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。

6 `9 {* v, Z& t' r# ~
语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

; |! P$ }7 s0 S0 r8 w  u! a+ r) nDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
+ J' \" i9 u/ v; s( YUperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
* ^. G( b/ h9 D7 o: o! W# }9 `SegNet 用于图像分段的深度卷积编码器-解码器架构。
4 ?( L1 C7 j) O& R$ y3 p) z6 K: b4 b还有（DUC，HDC）、PSPNet等。
' P' z3 |% X* F: a
常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

5 T0 ^3 S2 `1 T1 Y" w! d6 a对于损失函数，除了交叉熵误差，也可以用这些：
% }4 X9 j4 F* g
/ d, A( K7 {5 M: k* p, _! z5 RDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
/ E0 h8 O( c) P9 ]$ M: A7 Y) uCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
5 N7 j1 q) v8 ^  G) oFocal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
( g2 `" n7 V/ h) K6 g6 ]$ b

  G4 J9 k8 p- w$ G9 T2 }8 E
4 b: T) ]( ?$ Z
* Z% y+ d. j( k8 e$ Z% ~# r
( @4 f0 \5 ?4 h: r0 M————————————————
7 R6 y7 v# b5 M2 u' A版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
! x2 z- ]. l! z, e4 _' Z
/ Z3 N9 T6 o8 u: C9 o