Pytorch实战语义分割（VOC2012）

杨利霞

) h# u6 \" j1 |. }
0 y5 y9 w% E- D7 ~1 [1 B* bPytorch实战语义分割（VOC2012）
" N# S6 F7 n" ~: O本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。

! J* j, d) p# Y# [9 f8 g" c
+ d. J$ j$ j% u3 C$ Z1 X语义分割中图像有关狗、猫和背景的标签
文章目录
& A3 C2 O3 f' N! |& D; v6 m
1 图像分割和实例分割
  l0 k6 U6 I" S' K2 Pascal VOC2012语义分割数据集
2.1 导入模块
. ]6 H& ], e$ @# Z. [2.2 下载数据集
, U6 V; F% k' a0 B- C3 {2 ?7 ^2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
3.1 数据集类
3.2 读取数据集
4 构造模型
7 e3 M" J* E, W, H0 R0 R/ K4.1 预训练模型
4 ~1 t, I7 @5 e+ k, G, T4.2 修改成FCN
4.3 初始化转置卷积层
4 o" m* S1 q2 E5 B+ ^$ H6 G5 训练模型
) I1 w. |) T5 ?- _9 r8 @& O6 测试模型
' H) f7 g( H! o' v) a/ i) g( C& {6.1 通用型
6.2 不通用
! M! o) ~) ]# s5 v' \" O7 结语
& {- |( r- L) J2 U1 图像分割和实例分割
  q4 ~, @& b4 |5 x
计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
% q) B9 ?( D% e9 k; _7 R实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
& D7 |1 n$ L, f; u9 j, U2 r
; r9 S8 S$ |& j7 }! |- Y2 Pascal VOC2012语义分割数据集

2.1 导入模块
import time
( _( C7 k7 u' Wimport copy
4 o) }: e, G) m* n* Eimport torch
7 M; }$ M7 c5 s% [) j9 p9 ufrom torch import optim, nn
; G2 d: h& _  [# I5 Y- I% Q) Limport torch.nn.functional as F
! U6 D5 A9 m+ @& d' E5 |8 @, |7 Cimport torchvision
from torchvision import transforms
from torchvision.models import resnet18
1 d" z( K' {9 }. }. j* w% ~import numpy as np
/ I+ Q! I% g2 h3 |) |- Pfrom matplotlib import pyplot as plt
# g+ l' q1 I; G& I- f2 Ffrom PIL import Image
import sys
sys.path.append("..")
from IPython import display
3 L- C6 N* d* y5 j& Nfrom tqdm import tqdm
import warnings
$ d1 ~1 d- _, w( i+ pwarnings.filterwarnings("ignore")

' V# y, F: n/ c& W2.2 下载数据集
; f# i" @9 C0 s7 T8 c" l2 M
/ Y& r) \  }4 l. [1 K5 O" o, c语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
9 B6 t  v( r( p! b$ `: l

* k3 x6 T, h7 o5 LImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据
. b) p3 t0 I: r) M
9 c7 ^  h! M3 U& b定义read_voc_images函数将输入图像和标签读进内存。
& `/ v7 e& ?* |- s) y9 e5 g) X. y
. J3 g1 p3 _( N$ }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
    if max_num is not None:
# \- r; a( T3 S- U: C8 ~        images = images[:min(max_num, len(images))]
; G- a5 @1 ]! H: K' L$ p8 w7 c% K& ~6 \    features, labels = [None] * len(images), [None] * len(images)
* A6 [( g1 u9 P: f( b  w( b    for i, fname in tqdm(enumerate(images)):
8 S: g( A! q  L) \6 d5 L; y        # 读入数据并且转为RGB的 PIL image
        features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
' d  ]' k, O/ k9 q0 ?        labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
  @% i3 M2 u! a6 Y& S    return features, labels # PIL image 0-255

定义可视化数据集的函数show_images
  |- Y$ }4 Y# v1 O7 ~# k
0 d' S0 l& ]: P2 x# v1 @# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
' d- p, e. ~) \0 z6 A" G0 Q( E    """在jupyter使用svg显示"""
) p' V7 X. N/ F* R9 N5 C    display.set_matplotlib_formats('svg')
6 V% s! ]  }* n    # 设置图的尺寸
0 L" Z, j- ]4 Y/ i    plt.rcParams['figure.figsize'] = figsize

" z5 a9 y& G  I$ D+ ?) ]def show_images(imgs, num_rows, num_cols, scale=2):
/ W* j2 W4 i3 v4 f: v: i    # a_img = np.asarray(imgs)
# I0 m! D  C( D. T; j- {! i3 x& @4 }    figsize = (num_cols * scale, num_rows * scale)
! Q$ x: J0 \3 n( ?1 X0 B  J    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
% A5 Y% U2 F; x8 A: E2 E% X% g    for i in range(num_rows):
        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
. f9 w1 h& b7 @. \& B            axes[j].axes.get_xaxis().set_visible(False)
; [; |0 j( K: w' I" B            axes[j].axes.get_yaxis().set_visible(False)
% i9 B6 v# L9 m7 ^    plt.show()
    return axes
7 {% ^' {8 {- K. }
定义可视化数据集的函数show_images。
" F& t9 B, t* `& D6 ]9 e/ a
# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
4 S3 R3 T: a- Q& d7 w    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

def show_images(imgs, num_rows, num_cols, scale=2):
, B% [+ k8 h, d    # a_img = np.asarray(imgs)
! W6 R+ k3 g! q5 f0 j% I* u4 M    figsize = (num_cols * scale, num_rows * scale)
8 b% U3 W7 F* h9 i4 H4 ~; d    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
; v- H0 V! X2 ~5 J/ C8 n9 A        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
4 |* T+ b1 T( N0 I            axes[j].axes.get_xaxis().set_visible(False)
+ O6 N& s8 I# C8 S% H0 Z! x3 M            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。

# 根据自己存放数据集的路径修改voc_dir
voc_dir = r"[local]\VOCdevkit\VOC2012"
7 K/ w( C- d- M* j  u+ u0 @; S0 c/ }train_features, train_labels = read_voc_images(voc_dir, max_num=10)
- y! T- D5 _$ P" ln = 5 # 展示几张图像
( Z. {7 ^# C1 Mimgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
4 Q' b' U8 P2 c2 l( _% E. c$ L0 z
$ z& O8 b5 h0 Y- v. z$ c
8 i! M. ^2 o0 I' u3 s. ]
; S% ~# u0 w9 T  n3 d列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
3 H+ v9 g! M. a8 T- Q                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
2 w/ r# I$ s1 ]  \; ~                [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],
                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
2 d& D! k+ U+ Z1 m2 q& j               'diningtable', 'dog', 'horse', 'motorbike', 'person',
9 v- u, j& {; I4 l               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
* Z( p, F' R6 V$ h有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。

& x4 |$ a" o7 F5 v9 {+ C, L, g有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
for i, colormap in enumerate(VOC_COLORMAP):
    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
$ h# C% H- S; {: d. U  W; E& |
# 构造标签矩阵
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中计算的下标
( L. Y: ]! D$ c/ A
可以打印一下结果

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

2 J) C! j+ k$ t/ F0 C3 R2.4 预处理数据

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

1 q# v4 \5 P; ?  s" Y5 t
5 \: N, S' i4 W0 p  T" [3 odef voc_rand_crop(feature, label, height, width):
: U* D0 F$ n6 y, G& z. J4 f    """
2 a2 ?9 F3 j4 ]5 p5 p    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
+ g, |" m, {# {    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.
4 v' ^  q; n+ R+ V    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
) W& D# H! k$ {, e4 F) q. N, i    """
    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
( z! d* H1 d/ a    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
6 \/ w/ Y) L$ P$ f" a4 i" b    return feature, label
4 ?2 n9 V; Q" b1 f  c; q" q
. l3 O" V( C. z$ |6 N/ O8 b# 显示n张随机裁剪的图像和标签，前面的n是5
' [( J( ^! j, v% F" P. zimgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
; q! j. `  ]* ?) b1 bshow_images(imgs[::2] + imgs[1::2], 2, n);
! t0 h3 T/ e. h, T, p9 N7 A! T9 ^
0 r* m. Y' B2 K2 a2 ^

. Y9 j0 H3 V! R% P6 ~

' r5 ^7 D+ I, Q$ ]3 i
3 自定义数据集类

: t# \9 Y0 g- E! l7 w' n% v$ c3.1 数据集类

torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

) W" i& ^% k5 v( l8 E; z' _- [__len__ 实现 len(dataset) 返还数据集的尺寸。
) j) P1 \8 P+ |# e' Y9 h9 T__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

/ o9 H* f1 U1 Yclass VOCSegDataset(torch.utils.data.Dataset):
- q. B; M* R. ]    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
$ [1 }$ z0 H/ H/ q/ ~- U        """
        crop_size: (h, w)
2 c8 p$ a6 _5 J        """
! z: Y+ N; b3 R- r        # 对输入图像的RGB三个通道的值分别做标准化
. Z2 H% Q$ G7 k* o" b# h        self.rgb_mean = np.array([0.485, 0.456, 0.406])
        self.rgb_std = np.array([0.229, 0.224, 0.225])
0 k, @5 y1 l3 F) O& M        self.tsf = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
9 t5 A. g+ b6 ?% m( P7 L& A9 L        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
0 A+ j+ X% q7 R2 P% S# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
        self.features = self.filter(features) # PIL image
# ^9 Q$ D, P! Z9 G$ B$ T        self.labels = self.filter(labels)     # PIL image
        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')
" J. J7 `+ {: p; a5 ?0 P5 Z; Q
    def filter(self, imgs):
# E8 R- k  t4 S" `9 ]        return [img for img in imgs if (
. r9 Q7 }" d! G( c9 D3 C            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
: M/ D/ R. s7 M& l' I* s) W
6 ^' T" @/ d( |' b& ]6 u) r: [- C    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))

    def __len__(self):
        return len(self.features)
1 O7 ]" U7 [5 w3.2 读取数据集

通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

batch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
/ P2 v9 [! m$ T# O
: \( a% `, T2 C8 r# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
- l5 D/ T$ L9 g2 y& `- e# J3 N
- f& w- r" E: ?# A# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
num_workers = 0 if sys.platform.startswith('win32') else 4
5 k, S, F: M. I6 i$ otrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
+ b0 ^" W! @9 ~# d5 ftest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
4 Z4 ]( o; C0 [- E) w                             num_workers=num_workers)

) n' x# W( y2 r3 S# 方便封装，把训练集和验证集保存在dict里
( h1 T) D+ u% s$ J3 S" pdataloaders = {'train':train_iter, 'val':test_iter}
( L  v1 `+ @+ \% J" \4 w* Idataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}
% a+ a. c  r( g( R( A4 A
( `2 F4 g' x$ r, w2 z% H, F4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
$ K! U0 B) S6 @* D& Q
num_classes = 21 # 21分类，1个背景，20个物体
- i+ L1 ^  g+ t% g: Mmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

: l- @" E. y: h( g# 特征提取器
for param in model_ft.parameters():
    param.requires_grad = False
9 v+ b* n# ~8 o& {0 f# t4.2 修改成FCN
; x6 ^! ~1 E! `  y# S, X
全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
5 B. k: e  U) L+ L4 x+ m最终输出的通道包含了该空间位置像素的类别预测。

; G; ^4 v! [& g' I/ K对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
! ]. W2 Y+ t( L6 C
; m  _3 X1 s- ^" L4 V可以先打印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], # 去掉最后两层
* w! S" M0 j' Z7 [5 N( y              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
* }/ N" }/ p7 ~1 C( M! g              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小

9 t! U' t8 s! U: v/ [) C# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
2 B# v) j' h) W7 n. C/ p; K
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
; R5 ~* u6 w8 B! V) j8 `; u" D7 ~+ J#     print(X.dtype, X.shape)
% P! _+ g5 p2 h' N, K. L- x#     print(Y.dtype, Y.shape)
#     break


4.3 初始化转置卷积层
  j) Y; ~1 K- P% y
! x$ P  _3 j+ S: `* B在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
在坐标 (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函数构造的卷积核的转置卷积层来实现。
7 q. w- ]0 `3 r# v. j

8 i# S- e* q1 m: P5 s6 j3 |- F: ?# 双线性插值的上采样，用来初始化转置卷积层的卷积核
$ A5 f+ L/ b+ o( b: idef bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
( f7 O2 u0 s6 v, c* A# e        center = factor-1
1 B1 V% }0 Q5 w5 @7 t    else:
5 B# ], Z  {. i& y        center = factor-0.5
8 m& ]) |& _' f5 o    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
4 }( {' ?9 ^$ D6 L1 }3 `; `. `    weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
    weight.requires_grad = True
' U3 i4 A6 d$ p: f5 F    return weight
+ r4 V6 s$ z0 y9 Y2 w3 c
0 Q. o$ @/ \0 \# k6 T) E/ e
0 i- n/ I4 I. a& j在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
0 N# j! D2 F# w# q
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)
  Y- P4 a) S( b* g

) M% Q" E; B+ o  h& j+ |) T
5 训练模型

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

" d& z8 z6 {& n4 tdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
5 H/ O3 i: e# S# {- n( q3 a    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    # 每个epoch都有一个训练和验证阶段
4 `$ ~$ r/ i( d1 p7 E    for epoch in range(num_epochs):
0 {8 G7 u8 F% x6 V7 \- i        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
$ u( z+ ^4 I6 L% c7 Q' b        for phase in ['train', 'val']:
# `4 ~' b7 E& ^1 h) X, z) `            if phase == 'train':
                scheduler.step()
                model.train()
            else:
                model.eval()
1 }$ {  x; w+ Z1 N: V            runing_loss = 0.0
            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
3 k, h8 t% A" n$ F# ?% J3 I5 Z0 R4 ^                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
5 R; Z% {8 m5 x, R( y1 ~                                # 前向，只在训练时跟踪参数
  t1 P! |! k# e" A- ~                with torch.set_grad_enabled(phase=='train'):
. p" d: y- v6 g! o( H) @) Q                    logits = model(inputs)  # [5, 21, 320, 480]
+ Y8 C. ^6 B# t3 v                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
) _7 }% ?6 o- ]  ~                    if phase=='train':
                        loss.backward()
* q4 ?2 }6 Q, U0 A+ h1 \                        optimizer.step()
, I5 }$ J7 Q1 O! U: J% n& R                                # 统计loss和correct
& ?! m% Y# \5 t4 S6 I  I  a                runing_loss += loss.item()*inputs.size(0)
                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

  d- b$ y3 m+ K* _1 }            epoch_loss = runing_loss / dataset_sizes[phase]
            epoch_acc = runing_corrects.double() / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
3 f) Y- D3 S  H' O" ]! q                        # 深度复制model参数
            if phase=='val' and epoch_acc>best_acc:
                best_acc = epoch_acc
3 A0 f- q% b9 i4 n                best_model_wts = copy.deepcopy(model.state_dict())
        print()
    time_elapsed = time.time() - since;
) C9 z4 m; u# W4 h1 A0 ]3 p    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
; Y0 W0 {- ^7 r' c, k  W( l    # 加载最佳模型权重
1 S! M. \+ L2 N+ a1 I    model.load_state_dict(best_model_wts)
    return model

7 J, @7 e2 T) B+ ]下面定义train_model要用到的参数，开始训练
" k0 y' T' y9 E  i& _
) [1 @9 O* V6 g2 zepochs = 5 # 训练5个epoch
1 ?& A. L/ a7 O' m8 W+ _9 xcriteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# 每3个epochs衰减LR通过设置gamma=0.1
4 Z$ S2 c( H8 u5 E: f  L3 S$ cexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
& u" d& E3 O- Z0 k# y( s8 v
+ W* s& A6 D; C/ Q1 {; T3 d6 测试模型

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

def label2image(pred):
5 P: f. J3 B3 t( g+ p4 t    # pred: [320,480]
  U8 D/ |$ e8 J# ^9 e    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
- e: `( Z# O* d  p4 J2 A    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()
" v+ [: }: y/ S% q) ?( Y

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

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)
+ E' S  r8 V2 ?6 N# u( i% Qdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
, g2 F6 d& V( W    model.eval()
    images_so_far = 0
    n, imgs = num_images, []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
7 `' J* O+ [! t$ o            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 # 记得要变回去哦
* L* X" [* r0 `$ E9 U. b& ^9 \
& d5 |( E! s/ w3 }" f' P            for j in range(num_images):
                images_so_far += 1
4 Z5 O7 B- t* w; q- N                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)
& i* q7 P$ G* M8 Y# D5 X0 A                    # 我已经固定了每次只显示4张图了，大家可以自己修改
                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
# k4 \9 b+ W: s& s+ P) a                    return model.train(mode=was_training)
/ h3 n3 v" N# p# M% t/ D8 ]6 v  D
# 开始验证
visualize_model(model_ft)
- {8 N& }6 l$ a; q( D( X$ L6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
$ |9 A( W2 S7 q+ Ldef predict(img, model):
% O3 D* r7 W3 o# o- i; ?* V! b) q    tsf = transforms.Compose([
            transforms.ToTensor(), # 好像会自动转换channel
2 |) b5 O9 h8 t. \            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
- y4 ]( x' V9 X3 Z6 d5 M% G    x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
7 D/ Y# v( R7 F* p    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
% s' g* D5 ]& g7 |    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)

$ T! N. W0 h% D+ p- kdef evaluate(model:nn.Module):
    model.eval()
0 k# c% x. w5 H) J    test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10)
- O+ A- n$ C( l( J  F    n, imgs = 4, []
* P3 s3 U& |$ r/ c  [    for i in range(n):
9 ?1 u! u, L/ _/ l        xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image
        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
, Q4 {' O' f: C- s7 t- {    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)

# 开始测试
! N. |9 K/ H* w! nevaluate(model_ft)

& x& D7 f9 k2 b7 结语

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

4 f  w9 m4 K9 M3 X  c) \Epoch 0/2
----------
train Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456

* x9 Z5 P1 s( x! TEpoch 1/2
. i* A# Y1 v8 p. p- ~0 P' I----------
train Loss: 1.1288 Acc: 0.6535
val Loss: 0.9012 Acc: 0.6929
0 L0 C# d* |# X) W
Epoch 2/2
# T# t% ~' ?! U& K----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948
% ?( s, |, Y$ L# @; J7 M* K9 }4 g
6 Y2 g3 K8 F6 t( fTraining complete in 6m 37s

1 x3 A- Z" c/ @* |

. d. F; n- Z; |0 E0 p
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。

# Y( K& ^! J& T# Y* q* f
5 d5 g  b) b2 t7 O) E8 e语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：
* z8 j& s! x0 v3 f* W6 a
Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
# {, U0 C4 P2 z) e2 R& OGCN 通过全局卷积网络改进语义分割[论文]
4 F" W# b: Q7 s( F8 C9 x* }UperNet 统一感知解析
3 r7 F+ N1 J+ m2 ]9 [; U. NENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
/ \" _: z6 b2 K  u# f/ _还有（DUC，HDC）、PSPNet等。
5 b: s* @% b; l* j# K6 G/ ?1 D
3 m7 S" a. x6 k! t常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。

对于损失函数，除了交叉熵误差，也可以用这些：

" [4 y# a  e6 HDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
1 s# ]8 `% [% HCE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
3 G' L  v+ U! T1 s) M( kFocal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
: ~- J: |/ _4 R- s2 p7 E0 N8 PLovasz Softmax 查看论文：Lovasz - softmax损失。
' @  r% L+ }' H




————————————————
1 A' Z& ?4 L0 _: c6 ^' s1 ^版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507

8 Q& F" r, B2 {5 P1 T% ^# D( _

% b& ]& f  H" M$ x$ H