Pytorch实战语义分割（VOC2012）

杨利霞

# c1 |1 p' t+ G
Pytorch实战语义分割（VOC2012）
# X. J2 E! `, H" r本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
9 M8 ~" y4 h: A- R4 N: i8 t. s; D/ S

7 `  I' G8 [5 P5 H, f语义分割中图像有关狗、猫和背景的标签
1 z' Z; I9 ]( c' _& \3 `- O" B文章目录

1 图像分割和实例分割
* e* N& B7 V! c  [2 Pascal VOC2012语义分割数据集
2.1 导入模块
( T( m. X7 d5 R( B8 r2.2 下载数据集
2.3 可视化数据
2.4 预处理数据
* r! \. s2 R% M3 自定义数据集类
$ X/ J+ R0 X1 f& B6 f3.1 数据集类
' C- h( F" o7 j/ V3.2 读取数据集
3 r  h% O# q# ~1 x% f6 Q4 构造模型
4 G" E% u" m7 i* O* y4.1 预训练模型
/ k! k% H! k. t2 B( C, c5 g  Y4.2 修改成FCN
0 {, n( }. D! ^# |4.3 初始化转置卷积层
- B; Q8 i/ o* S2 h4 D' [5 训练模型
6 测试模型
$ z' k( e9 Q% N/ K$ Y" N6 w6.1 通用型
6.2 不通用
. F9 L6 S- l+ A! I9 |* U7 结语
1 图像分割和实例分割

, N7 N4 B2 d3 |) m计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
  Z: j: j' k- w, ~% x5 V1 I
0 X& u( v" j) ~图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
  V8 O- a9 a; X& Y! r5 \) t
& B9 J$ A2 J. i2 Pascal VOC2012语义分割数据集

2.1 导入模块
import time
; C& r3 J. [6 {' X; Cimport copy
import torch
/ W2 R5 m0 o- w/ m, D% vfrom torch import optim, nn
9 E! s& q6 y  L9 |import torch.nn.functional as F
import torchvision
from torchvision import transforms
" X3 e. l& F! @* ]& r& a7 K- S9 `from torchvision.models import resnet18
import numpy as np
4 ^2 ]+ q3 \- x7 I7 Tfrom matplotlib import pyplot as plt
from PIL import Image
import sys
sys.path.append("..")
- g9 K: H) i) }) dfrom IPython import display
from tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

2.2 下载数据集

语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
5 O) [1 u1 a, j( ], J6 e  f/ w
' g! }  E: s* w. [4 g2 K; N' ?$ |
ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
; f+ ^( J! ?; V6 b; x) ?JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
0 M( _* B* y/ s# Y' a3 k; J2.3 可视化数据

8 Y5 ~3 F: n$ i. W8 y定义read_voc_images函数将输入图像和标签读进内存。
& U2 K  X) N1 P6 o7 |
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')
* ^+ O) b2 g; p+ S. B# T$ O' S    with open(txt_fname, 'r') as f:
8 O. Z0 l/ X7 l$ w1 j/ V        images = f.read().split() # 拆分成一个个名字组成list
) y/ ^* [1 Q3 \$ b3 E" U    if max_num is not None:
) D# r+ O! `( J3 s" t. Q        images = images[:min(max_num, len(images))]
! {2 L/ K9 w, i# A, k    features, labels = [None] * len(images), [None] * len(images)
/ C1 |9 ]8 O6 w9 O$ v    for i, fname in tqdm(enumerate(images)):
8 D" I' U/ D# @) D- t( S3 F$ W        # 读入数据并且转为RGB的 PIL image
: |* `( b* L. m4 F; P        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
6 K# F- d( z* A) b        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
    return features, labels # PIL image 0-255
+ y; u0 @& G/ n7 @& z& t
3 [' Z/ q, {/ J$ k7 N3 u. d定义可视化数据集的函数show_images

% e# Q$ \  ~* R: k) o# 这个函数可以不需要
3 D5 Q: N" v" a8 E  n& s2 tdef set_figsize(figsize=(3.5, 2.5)):
; }! B" t/ x; ^' f+ }# O    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
( g% W) b) Q6 T7 f8 W/ D    # 设置图的尺寸
    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):
% {9 b+ G0 ]3 y  c1 N        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
$ H) n, P5 A  ?" W/ f            axes[j].axes.get_xaxis().set_visible(False)
$ N3 h3 @* P2 o1 G% g            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
2 E* T/ i# Q0 b* C5 F: I    return axes

0 M0 i% S: a: x定义可视化数据集的函数show_images。

4 B, _& F% [  \" e) B# 这个函数可以不需要
+ ]" k5 \/ O7 U) D6 ?def set_figsize(figsize=(3.5, 2.5)):
4 [: f" K% j( t  Y    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
; _1 D( g9 z& V- w+ F/ Z    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
& B! p# d( L+ S, \- v    figsize = (num_cols * scale, num_rows * scale)
! V. Q& K1 z" o( Y; u) p    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
; h/ f2 E6 s0 G( F" `    for i in range(num_rows):
        for j in range(num_cols):
6 r0 `8 y- o& v# c- ^( }1 i& ]4 H            axes[j].imshow(imgs[i * num_cols + j])
1 Z  o( a: z2 q8 H* J, `8 q. U            axes[j].axes.get_xaxis().set_visible(False)
2 }) J7 R3 t  u( _            axes[j].axes.get_yaxis().set_visible(False)
. t7 w( Q8 i: b% v    plt.show()
0 A3 B- Y7 e9 D: I5 Y    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

7 _. z4 v7 A! P- k: y* j7 S5 Y# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
- [  `8 Z5 U6 Q8 ]& dn = 5 # 展示几张图像
/ M) ]: ~2 ^% Mimgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
; {2 w/ T8 t6 u' W: O/ Q4 l
7 q; \; B) i, c! ^

4 g/ f3 n( q/ x2 E: E列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个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],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
! @& a) G- _3 d: h9 _                [0, 64, 128]]
- O* K1 Y# v1 ^1 y- V) }" j- Z# 标签其标注的类别
) e  z4 M8 i6 b6 j* `; dVOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
# j& W6 F2 \' y1 s4 S有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
) L- _/ j' v" o7 M- B$ G2 T& l+ i4 i
/ G: ~* H6 `- C+ Y( h1 \; D有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
4 V$ _0 {: ~3 F, Q' s+ _6 m0 @colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
" F; X, ?% Q# A6 c5 Xfor i, colormap in enumerate(VOC_COLORMAP):
- a. d  ~' n9 h2 b# T  j/ `    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
9 X5 R  i3 h7 H( y
# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
' x" C6 N2 S$ u* X; ?; a    colormap = np.array(colormap.convert("RGB")).astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
) x) Q9 Z# [+ z( v7 x    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

/ D: ?0 n* I4 H1 t可以打印一下结果

5 Q$ ~* T$ R4 V! Yy = voc_label_indices(train_labels[0], colormap2label)
% O3 u% G6 _4 \7 T* E* kprint(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

4 B$ b+ S( |+ F) h% G, |; x5 ^2.4 预处理数据

在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。

8 Q- @& \* ]9 M# z4 H5 T) b$ }! S
0 g5 t  y! S5 \def voc_rand_crop(feature, label, height, width):
3 h" ~9 @8 k+ F! m6 ~3 _5 {    """
  M6 m- ~0 Q; Z8 J  Q1 _    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
" V. ~/ i8 v  O  }    Get parameters for ``crop`` for a random crop.
1 E% l( c% @; f+ _    Args:
        img (PIL Image): Image to be cropped.
' b7 l& ?- [: o  ^* w0 V( y        output_size (tuple): Expected output size of the crop.
    Returns:
) f5 v" |/ Y6 j6 D( Y& |        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))
: M8 T* ?, M( g% Z* E7 _" _0 z4 a    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
    return feature, label

# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
for _ in range(n):
* M! m# s9 k1 ^  P    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);




8 N/ y% W1 e" g: G
% {: i% e9 y8 g, q  o+ Y; O
3 自定义数据集类
- B8 W: g$ O6 r  e& g* V2 H/ g* M1 l
3.1 数据集类
  C5 k: K" f5 s* s1 ?
torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

__len__ 实现 len(dataset) 返还数据集的尺寸。
# o) n9 o2 C  y5 c/ p__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
' \1 c+ a8 R+ {; c. ?- ~由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

class VOCSegDataset(torch.utils.data.Dataset):
$ n8 {( u; x, G8 {    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
        crop_size: (h, w)
        """
0 H$ i" f3 _! E; l1 M* M        # 对输入图像的RGB三个通道的值分别做标准化
        self.rgb_mean = np.array([0.485, 0.456, 0.406])
) S) _9 |2 M- N- P; l7 [        self.rgb_std = np.array([0.229, 0.224, 0.225])
. e8 e& B0 M  Y1 Z8 t        self.tsf = torchvision.transforms.Compose([
$ e" n- ?, K# F  h* u$ }( n            torchvision.transforms.ToTensor(),
6 ]- t1 p3 d$ `            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
9 f8 T8 T; N2 _' v1 X' }        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
        self.labels = self.filter(labels)     # PIL image
" P# ~+ C! k  S: A, R        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

# Q+ E  @( o( F2 b$ \    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 w# w. y) s+ }4 @
6 P/ z6 k1 N( o0 n0 J" _- F    def __getitem__(self, idx):
8 Z1 O0 n; n4 I* w& w' E/ {( D        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
; y  Y4 H2 k2 y7 T- l( ?6 z! M4 }                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
! r+ E* o1 B. [0 L/ I2 t' n- Z+ e8 F/ A
9 V: j0 H8 n% G5 l    def __len__(self):
" z* B' }: ~, y* l' j6 \7 @        return len(self.features)
9 g" F" _  n7 m( Z0 G3.2 读取数据集
6 u/ ]- o1 q8 w
通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。
) A% ]3 d: i6 \1 X4 h. R! @
0 `: v; v+ L  f( x! l0 X/ i0 |2 [" xbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
/ t3 |; H3 Y4 L+ t, p6 H2 N
# 设批量大小为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,
, N( O: X3 w- K! X% r7 r+ G" i4 F; v                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
; A7 F# o! a: }2 N# D; F/ ]$ h7 Ndataloaders = {'train':train_iter, 'val':test_iter}
dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}

4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
, T3 c* U+ a( O" T3 `; e# c+ g
num_classes = 21 # 21分类，1个背景，20个物体
& r# _! E) X; T# U+ K! Fmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

$ S+ u0 N* Y) S- v7 U0 E: t# 特征提取器
for param in model_ft.parameters():
! C( r  M+ j  d# k. U    param.requires_grad = False
7 D! w, @& F; A1 K# y  t4.2 修改成FCN

& K1 r( U0 e, n% T8 }$ R全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
& I. W; c  q9 C/ k# N最终输出的通道包含了该空间位置像素的类别预测。
3 ?% B# e# B2 w4 h1 D
# Z& F2 h$ I  d, H/ z: U' c% {对于转置卷积层，如果步幅为 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 倍即可。
+ M8 R" j3 c8 O
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
2 P4 r; s; l* B, ?- T  }              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
' L5 x. S4 j- K# U              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
* l8 H. g% E) h, J- y5 K
# 对model_ft做一个测试
, a# j# R7 Y0 n* X1 Tx = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
5 }6 _8 _% }  J- p  J#     print(X.dtype, X.shape)
4 m7 p8 P/ n* T" p5 N#     print(Y.dtype, Y.shape)
#     break
3 b) G, O( h4 o- E; k
* `3 I# h, g3 Z% s4 y
4.3 初始化转置卷积层
; \! B! n  e" r' B& M
在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
, O& t* l' c2 j) S+ l) g8 x在坐标 (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函数构造的卷积核的转置卷积层来实现。
% Z, l+ e1 F' Z! k0 c

; k$ t# r/ l" v7 k( i' s8 I& s; c# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
( ~6 }) |: C+ C" s5 C8 r    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
        center = factor-1
    else:
        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
7 |. M! k9 `# H  [, H    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
( B1 @- K5 U& L( e    weight[range(in_channels), range(out_channels), :, :] = filt
, R5 S2 e& p! J  T/ K    weight = torch.Tensor(weight)
1 |+ L8 {) I, e    weight.requires_grad = True
+ c! D: S5 H8 @    return weight

7 z) Q8 t: T% A2 [6 x; j  i
在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
* J) {0 _9 J3 K* d& K0 N3 u
0 r- i# U; I- snn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)
model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)


) c0 l6 q/ a# Z" p& G! G
5 训练模型

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

def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
, e/ k9 R) P, r5 j    since = time.time()
# z6 S! y! @) f4 P6 A- V, u    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
  V! {4 `1 f3 n* Q' G! `2 x        for phase in ['train', 'val']:
1 [2 T3 G3 l5 ~, {" O) ?            if phase == 'train':
                scheduler.step()
                model.train()
            else:
8 G( x" y6 s7 {, n                model.eval()
6 N) h9 h$ M7 C. B            runing_loss = 0.0
& v# C% _/ z  ?) U# \0 v  S            runing_corrects = 0.0
$ \3 r+ i' I. [  w" m/ s            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
% [" x# w# K9 h: j; i, L                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
- E! D$ s6 o7 r                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
                    logits = model(inputs)  # [5, 21, 320, 480]
; r& `& s- L' `- H( W/ J                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
                    if phase=='train':
7 `0 t' ?: Y: @8 Q! w) S' a+ C5 ^                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
1 z9 G2 H; a- U* ?                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]
# K4 Y4 w$ C; }( A2 D- n2 P0 e5 l1 q            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
1 U1 Q. y" Y) q( d  z                        # 深度复制model参数
7 \; B+ J) ~7 b& S8 ?, n6 a            if phase=='val' and epoch_acc>best_acc:
, w% @8 [/ S, ]4 @6 s( `4 Q  I                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
        print()
7 i4 h2 d' D3 ^9 f1 j- A( C2 N    time_elapsed = time.time() - since;
$ \1 K" R" a  N; |* b- ~7 p    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    # 加载最佳模型权重
6 ~7 J9 {$ C% M1 O/ d3 L% Q4 k- B    model.load_state_dict(best_model_wts)
    return model

( j' T& _' h# H& E, b; _, h下面定义train_model要用到的参数，开始训练
$ D5 e9 h- w2 q) Y8 l
epochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
+ S$ \- ~0 f& @  ]/ Q( Voptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
7 d: x. `) O  u- G! d# 每3个epochs衰减LR通过设置gamma=0.1
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
* r% P" x$ _: G4 M& P
# 开始训练
+ u' B5 D7 `" \1 U8 B& mmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
4 N4 G1 K/ A# R/ f* a# e
; [& C" W. U& @4 G8 m6 测试模型

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

def label2image(pred):
    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
* r# n! ?. R7 x- B( Hstd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
- y: [. @- B7 j# f  Jdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
    model.eval()
7 b8 Z- |  `8 p4 T. D# E    images_so_far = 0
    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)
6 {1 Z2 I* H/ f8 \8 l9 J            pred = torch.argmax(outputs, dim=1) # [b,320,480]
; ~8 K! n( q8 n            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
& S: ^# m5 `* Y
6 c4 F) N3 q% P8 x2 C            for j in range(num_images):
                images_so_far += 1
                pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
                imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]
                if images_so_far == num_images:
                    model.train(mode=was_training)
. c0 w  g; `8 k5 u. W7 \% M' r7 y                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
6 h& Z, Y: S! Y, l9 i' H6 \                    return model.train(mode=was_training)

# 开始验证
- b) S: K; `9 K5 p, Nvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
def predict(img, model):
    tsf = transforms.Compose([
            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)
$ z- L8 u* n& J% o. ]( ^* v. ~    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
- a7 o; x. I" T  R6 j
+ S0 @' m5 d4 i5 _9 J, J! Edef evaluate(model:nn.Module):
5 J( G- P* q/ G% J    model.eval()
! G: i" m- i3 a' t0 v! Y+ S) \    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))
        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

$ \, Q) ]+ Q# l, x* m7 }# 开始测试
evaluate(model_ft)

! j9 S& v, M+ t! B& z7 结语

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

8 t" T7 ]: u' d6 |6 HEpoch 0/2
----------
train Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456
5 @, ~( f& t/ u* D( r+ u- c
  p/ }- K% Q) I7 ~Epoch 1/2
8 X: t' |: j: h) x7 I----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929

Epoch 2/2
----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948
7 v. b) E3 R2 O  k% D
4 s) H5 ~+ l! J' ]Training complete in 6m 37s
% `- z+ \5 P' R3 E' n


. p, |4 k: s- W' m. R& k8 ]
& L( K8 s& E: t% Z当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
* f+ b: |" |6 }/ i/ T+ x2 Y; y
# H2 O" Z" Z7 c0 W* i# W8 R
语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
7 i) ]2 V" J- r7 h' f! W
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
/ c. c# s% X7 s+ TU-Net 用于生物医学图像分割的卷积网络
9 i& ^) m. q& K" i0 R1 s& YSegNet 用于图像分段的深度卷积编码器-解码器架构。
还有（DUC，HDC）、PSPNet等。
6 [1 R! i) z% T5 k8 t; r
" h3 r2 e1 S# r; Q1 t常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

3 v& M* ^% e6 Z' h  ?2 J' k对于损失函数，除了交叉熵误差，也可以用这些：

Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
! C* H) r$ D$ uCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
5 ]! D3 g& R$ z: ?. ^Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
Lovasz Softmax 查看论文：Lovasz - softmax损失。
7 U$ ?0 j7 _  ^& w
1 `" z! W& n8 U' x+ ]8 Y3 i
6 }0 f- O% u2 ?6 L
0 r: m* D6 c4 a5 d0 O, V: A; Q3 d
0 x! I; k5 x! ^$ k; z& o6 ]; c
————————————————
8 y- O  |! J- U# j6 U# s版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
1 ^5 Y6 _2 Z$ |原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507

3 V" b! O' _. X+ H$ _& }9 I

4 H) F7 O6 H7 g$ V: k( S# ~$ K