Pytorch实战语义分割（VOC2012）

杨利霞

; X  ?! o+ v5 W1 s# z  |
Pytorch实战语义分割（VOC2012）
本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
+ N  a# ^4 q7 t3 q' \语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
/ Y3 K& M9 I( Y! W. O, S& R
, N" S( A  d8 U4 z3 }7 g/ O
语义分割中图像有关狗、猫和背景的标签
文章目录

% I& R. h" b* X" A# x1 图像分割和实例分割
& g+ k- m# g1 @) O+ b/ l2 Pascal VOC2012语义分割数据集
2.1 导入模块
2.2 下载数据集
6 W4 @# p& q) ]  ^# B; L* D2.3 可视化数据
2.4 预处理数据
3 R. |1 m7 R( x7 d& E  Z3 自定义数据集类
) b# g9 c& |* O; N3.1 数据集类
3.2 读取数据集
! y) U, i0 l9 _/ {! r. G( w4 构造模型
4 {3 }1 d  q! u6 t' f- s1 b4.1 预训练模型
. ~/ V" x% B& s* m- O- Q4.2 修改成FCN
4.3 初始化转置卷积层
# ~* T+ h" l) ^; `' i5 训练模型
9 {  G, x, P, E) i3 e" b6 测试模型
9 k. k& ~, Z) P6.1 通用型
6.2 不通用
# k$ }7 ^8 W. n& d, i2 W6 Q7 结语
1 图像分割和实例分割
! d# s$ Q* B! b/ I2 g8 H9 o6 G
计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：

3 ^4 T, Y4 y1 c. e% n* C图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。

8 e. m7 l  d& l& w+ p" p' V2 Pascal VOC2012语义分割数据集

) g3 c* O% A( k2.1 导入模块
import time
" R" _! X+ \; R6 |, q5 V! a3 `import copy
4 v( W3 ?9 X: x* n+ C5 R2 t3 E6 Aimport torch
from torch import optim, nn
! F6 E5 e& z- {$ n: w- X, Yimport torch.nn.functional as F
import torchvision
: L1 j7 A9 o; E) L6 c7 J6 Gfrom torchvision import transforms
" k  S+ g# U- W1 Y2 s0 T6 D! K' efrom torchvision.models import resnet18
) G. B* V+ E, c; Yimport numpy as np
from matplotlib import pyplot as plt
from PIL import Image
import sys
& D+ h$ W( j* \( q$ ]6 s; Rsys.path.append("..")
+ j9 w# U; O! i2 Q' mfrom IPython import display
' S* w: R( d6 r4 Gfrom tqdm import tqdm
" W. [/ Z$ j+ F% n. V  `import warnings
warnings.filterwarnings("ignore")
% N% B* q# n) w: Q( W
! J7 |3 i1 n; K) N6 G* X' k) ~2.2 下载数据集

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

ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
4 k! q( r; k) v0 `JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
2.3 可视化数据
* _9 g2 R; v4 q& y* A$ `
+ n- o1 R  o& _+ ~定义read_voc_images函数将输入图像和标签读进内存。

) Q9 a: e- [/ O7 L7 e7 i. Ndef 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:
7 H4 l" ]1 P3 Q7 g* R7 |        images = images[:min(max_num, len(images))]
% x4 B1 Q. d6 Y/ w3 x; ^    features, labels = [None] * len(images), [None] * len(images)
    for i, fname in tqdm(enumerate(images)):
3 n& A/ o7 {. L" o. ~        # 读入数据并且转为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

1 g. @5 ]& B8 |: |$ W: \  p0 {) s定义可视化数据集的函数show_images
% ~7 }& [+ s9 _/ M% N4 t
$ T( B2 r1 F6 j& K# 这个函数可以不需要
5 o3 `! c) C: ?% fdef set_figsize(figsize=(3.5, 2.5)):
6 j+ j# ]; ?1 Z% y* c  E    """在jupyter使用svg显示"""
9 H+ s: u0 s- w# ?5 ^8 X6 ]; H    display.set_matplotlib_formats('svg')
! [* a" [* Q. U1 d! S0 ~0 w9 M    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize
  f$ B0 w2 N  Q/ O; ]2 W6 c
* c$ q8 O7 m1 ?. Mdef show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
. I! w6 j. t7 G+ K& p& n    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
2 T6 a& [1 N6 {" Y    for i in range(num_rows):
        for j in range(num_cols):
: f( _& Y' w  D% D4 `" E            axes[j].imshow(imgs[i * num_cols + j])
$ ~2 l  s" x. `) |9 }7 c8 S/ r& ]            axes[j].axes.get_xaxis().set_visible(False)
& Z+ V' F5 J- ~            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes

定义可视化数据集的函数show_images。

# 这个函数可以不需要
def set_figsize(figsize=(3.5, 2.5)):
$ L  K1 r$ `3 a    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
# y  y% n8 _+ j0 r/ C* n+ i$ d( L    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize

" W$ F" K) s& B7 o" q9 z) Fdef show_images(imgs, num_rows, num_cols, scale=2):
3 k1 I6 A4 X3 }2 Z2 y& f  D1 ^    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
: F) @. L4 @, v- \9 A+ F7 ~3 N0 [( o    for i in range(num_rows):
        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
* H  }; V7 ^+ P1 k- M            axes[j].axes.get_xaxis().set_visible(False)
2 M% y) P& J5 ]- ^& v) }4 p3 R            axes[j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
5 U, t* X: @9 }7 o2 X
# 根据自己存放数据集的路径修改voc_dir
2 C( |/ G2 e& o. x+ [: }voc_dir = r"[local]\VOCdevkit\VOC2012"
- T$ y# d, O1 Ktrain_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
imgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
* d$ w7 U- V/ Y- r

6 {3 t6 z7 J7 `( ]! w! Z* Y/ K( w
列出标签中每个RGB颜色的值及其标注的类别。
# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
. s! j, s# }9 ~2 B0 V                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
& T7 R2 P! r* R6 Z5 L                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
0 f9 q2 P; X: t                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
' X/ s9 m, E- h: ]- s- v- F                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
$ Z9 E( _3 k  V6 c& u                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
$ F/ R1 E4 Q' Y- l8 u有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
% q$ l, ]4 `. A: p7 O% }
+ ?1 [9 u% _. s3 B3 A! O4 v' A有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
) D+ [7 W7 c" ffor i, colormap in enumerate(VOC_COLORMAP):
9 T# l7 h6 q% |9 o. c    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
& o( s& E) v4 R% ?* u
# 构造标签矩阵
def voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
: d, t6 z4 I; _7 Y3 ^4 c2 I    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
. s" C; Z$ Z) D$ G8 |, {( K2 v0 d. C
可以打印一下结果
1 |  ]- @3 J; u; D9 m% y8 o
' v/ P0 F2 |' B, _9 Sy = voc_label_indices(train_labels[0], colormap2label)
5 }9 X7 }# b2 u9 dprint(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES
- `' M6 g3 A7 N) W7 v+ a0 K- H: F
2.4 预处理数据

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

' O) o4 m" n# f
def voc_rand_crop(feature, label, height, width):
    """
    随机裁剪feature(PIL image) 和 label(PIL image).
    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
9 K3 j9 w  ~7 {6 Z! A    Get parameters for ``crop`` for a random crop.
5 t, a6 w1 W5 l$ [: u( Z2 ?1 Y    Args:
        img (PIL Image): Image to be cropped.
        output_size (tuple): Expected output size of the crop.
! x1 ]  c) {9 H% Y0 J! n' e  {    Returns:
7 y+ B5 R4 x" i4 e: ~$ i" r: }        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
$ ~3 D, i! p. u# H1 l  ~    """
9 U% f& V: J( [& H- P    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
! u& S) u6 f, i# C9 r    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
. m0 V3 Q5 k: [4 {( F. T* l    return feature, label

# 显示n张随机裁剪的图像和标签，前面的n是5
imgs = []
0 H( V7 V5 D  Z& N/ wfor _ in range(n):
& x5 Z5 c3 W4 ~    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs[::2] + imgs[1::2], 2, n);
" f- {9 y/ N7 c! M/ A. z! s





0 b' X" Y) Y( w9 D3 自定义数据集类
* m7 R5 ~4 r3 |2 f. K- V% e
: z; m" J: ]0 ~  @3.1 数据集类

: ^. I0 l" p8 r7 \! y$ utorch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法

__len__ 实现 len(dataset) 返还数据集的尺寸。
. Y4 ^$ x: J4 _" Z__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。

class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
        """
7 T  [# k+ w: `' \- |- Y- b  ^1 _        crop_size: (h, w)
! u7 D( n' [, _' R  k        """
        # 对输入图像的RGB三个通道的值分别做标准化
7 C, c- O% Q7 d$ z        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([
" m$ ^4 D/ c6 n! ?- ~" q5 b            torchvision.transforms.ToTensor(),
1 j8 q: c5 I7 Y4 Q& P/ K. u) z            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
% `" f5 v* X: D7 y$ z, |        self.crop_size = crop_size # (h, w)
        features, labels = read_voc_images(root=voc_dir, is_train=is_train,  max_num=max_num)
# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除
8 ^: T( H( F4 n* K" Q' L4 w) n% Y* b        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
' [+ m  l7 u* z        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' valid examples')

    def filter(self, imgs):
: \# F5 L, W7 u8 X        return [img for img in imgs if (
            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
# I5 q  h% i" ^7 i3 X
! }6 E$ c. D* O, ]! u% x2 b$ d    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))
, b  t/ K0 T9 A' i1 w
& l( q8 B" ?% |    def __len__(self):
        return len(self.features)
3.2 读取数据集
5 {9 K, h: u) V. ^
通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络，所以随机裁剪的输出图像的形状可以自己指定，这里指定为320×480​ 320\times 480​320×480​。

8 r0 }# H* k, d- Y) m, G0 i6 V) vbatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
crop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
# C7 W, o, C  n$ fmax_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~

- A% j* N5 d7 n1 X& y# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
! ]4 f: Q* `0 D1 l2 A6 L( `6 A- avoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
4 x3 S* p: x$ C" E& d! ]6 S9 o' X
+ t. ]  N) r$ ?& P0 X) H4 _# 设批量大小为32，分别定义【训练集】和【测试集】的数据迭代器
3 o! z6 E$ f% C" b. C% {7 V; {# p, ~num_workers = 0 if sys.platform.startswith('win32') else 4
. e3 _( r  c- b1 p& ftrain_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers)
9 V$ c# m. w: p; C6 ]5 b+ Q$ _test_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
6 g) G. m( _! R& x' n0 n                             num_workers=num_workers)

# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
; c5 P; B! a9 b/ N" x- udataset_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')

num_classes = 21 # 21分类，1个背景，20个物体
8 n% d! _/ W, U! vmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数

# 特征提取器
for param in model_ft.parameters():
4 `9 Q! e; D2 f/ o: V' y    param.requires_grad = False
9 k4 v/ ~4 T# {0 h4.2 修改成FCN

3 o8 P- R# P9 [/ E( \4 V1 g9 z全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
) B5 ?. P- c% E( l, D最终输出的通道包含了该空间位置像素的类别预测。
2 o0 z0 [) Y$ `4 U: \0 Y4 `+ L
) M, q; g' T' Z# o对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
/ w+ \4 m, w" d
可以先打印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 倍即可。
8 X7 x1 ]1 n0 H
3 m8 P# c: E$ P  m. imodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
5 ~0 U7 M2 h: g9 a              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
" L( M+ s( O& T; }8 Z! m              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
( A# c( s& o$ C, \. g
. Y1 e# z" l. i. M7 U# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
4 l& J' t3 L' i2 n1 K3 V+ k- R
) E$ C3 S( A% K9 d% ]" ^# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
#     print(X.dtype, X.shape)
  \) H' ~) T0 e#     print(Y.dtype, Y.shape)
#     break
+ \0 [' i9 w2 E

, M: e9 o& s! w% f* K# y4.3 初始化转置卷积层

在图像处理中，我们有时需要将图像放⼤，即上采样（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函数构造的卷积核的转置卷积层来实现。


: g0 K4 ]) {1 }8 U$ j) c+ U# 双线性插值的上采样，用来初始化转置卷积层的卷积核
. X- \# X7 c/ u6 ]" ndef bilinear_kernel(in_channels, out_channels, kernel_size):
/ F, c0 q' z* _    factor = (kernel_size+1)//2
    if kernel_size%2 == 1:
        center = factor-1
    else:
* d7 s7 J# b( P/ P        center = factor-0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)
0 O8 M! U2 _% T+ K; X/ S4 L    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)
/ ?1 y" J! ^( c# m! C    weight.requires_grad = True
5 @5 X# P; r3 D, t3 p    return weight
6 |9 c6 O1 r' w2 v* f
$ ?7 {, ]$ m4 R( D% d
, b1 R7 r" T  u- t% M! I9 }在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。
1 n+ |0 c# Z! [9 _0 D  I6 Y2 U( L
  c7 J- V) R9 c+ 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)
( q* D4 F1 ^* s
, ^6 p( o- L/ I3 c5 @9 Z- v
) U/ w9 ~7 I0 E/ T% w
# O1 U$ N4 ^8 b1 c5 训练模型

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

$ s9 g% C$ a1 X$ H- {+ ^. C, n4 Bdef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
1 P" `5 z* v% k4 R9 ~1 q    since = time.time()
4 b) q; F' j3 }2 U9 E! C    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
* U1 ^5 }% r: j) v2 w: j8 \    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
3 I, I& w6 x+ |8 u+ [9 B  J" {: p        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
: b* W- F) k( Q9 b! q4 ?        for phase in ['train', 'val']:
1 a+ {8 ?7 }4 G% H& p            if phase == 'train':
# o3 N* N8 }1 U( a3 x. `                scheduler.step()
                model.train()
, K& B  @- e0 q2 a1 J            else:
* J7 O6 H$ B2 v9 O* X; F! S" u                model.eval()
            runing_loss = 0.0
" u, f! Z2 n: k" _! e            runing_corrects = 0.0
            # 迭代一个epoch
            for inputs, labels in dataloaders[phase]:
" F% z( P; b3 t/ A                inputs, labels = inputs.to(device), labels.to(device)
                optimizer.zero_grad() # 零参数梯度
9 J' ^) C2 O! C  d1 v. S+ n  t) Y                                # 前向，只在训练时跟踪参数
                with torch.set_grad_enabled(phase=='train'):
! I# v4 {% Y9 h# q                    logits = model(inputs)  # [5, 21, 320, 480]
                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
                    if phase=='train':
                        loss.backward()
                        optimizer.step()
                                # 统计loss和correct
6 i  q; o$ z. Q/ m4 y/ r9 _                runing_loss += loss.item()*inputs.size(0)
7 o+ ^  `9 ^5 ?* w7 y                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]
, a. t& Z5 k- g; _4 a8 Z            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
                        # 深度复制model参数
8 g% F: k  e2 ?. P; m2 e' \9 a            if phase=='val' and epoch_acc>best_acc:
/ d& U3 b! ?# Q                best_acc = epoch_acc
: G* Q7 Q3 |* z                best_model_wts = copy.deepcopy(model.state_dict())
        print()
3 ?: e2 _+ l& I5 R& {& l    time_elapsed = time.time() - since;
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
8 m8 V+ p$ [# d+ z9 W& O    # 加载最佳模型权重
4 C0 H6 p: r( f5 }    model.load_state_dict(best_model_wts)
    return model

! o0 m0 {- `( W/ i0 k下面定义train_model要用到的参数，开始训练
1 u. Z  f% P, G. D' G6 a% u) N
+ J# K" s" Z# B/ m- gepochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
# r4 _0 I2 U2 c/ d" ?# 每3个epochs衰减LR通过设置gamma=0.1
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

* {% ~1 D4 z" |% V3 \" G# 开始训练
! \! {, p2 x3 O9 N& Rmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)
( |8 R$ t; s. F# q
6 测试模型

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

def label2image(pred):
    # pred: [320,480]
    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
    x = pred.long()
6 t$ Z  f- ?2 O" `' |6 S    return (colormap[x,:]).data.cpu().numpy()

/ H0 x+ {; B* y2 |: m7 w

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
std=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
% y- }+ o/ o3 _! tdef visualize_model(model:nn.Module, num_images=4):
    was_training = model.training
    model.eval()
( c( I) `9 C3 {) Q9 F6 c    images_so_far = 0
4 T0 r. C5 m8 D; z: b; J    n, imgs = num_images, []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
( U+ D! i7 y8 S, ~! V) i1 B. z            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
            outputs = model(inputs)
" J* }; C; q3 X7 t* G# F" S8 P4 \            pred = torch.argmax(outputs, dim=1) # [b,320,480]
8 R( p* o" L/ i; l  c            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦

# d7 H' k3 h1 e3 O' {3 i            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])]
$ @% t* G* K+ ]. [4 ]                if images_so_far == num_images:
( Q1 ~) H  r. e6 U- X                    model.train(mode=was_training)
                    # 我已经固定了每次只显示4张图了，大家可以自己修改
' e+ X2 m' Q- g- I$ o                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
$ |! t/ G+ N# P; Z                    return model.train(mode=was_training)

! S9 I; m( M+ w9 @/ z; X# 开始验证
visualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
6 C. X+ A- c1 a: }5 x2 Cdef predict(img, model):
# K0 |- j4 Y! P) h1 Y* p  C    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)
7 {) E- G/ t( b) E* P# F    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):
6 s% Z3 e. w; ?9 i4 l    model.eval()
    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
: @) b, O* ~2 x" m9 b- W        pred = label2image(predict(xi, model))
        imgs += [xi, pred, yi]
: \3 ~3 [; j3 z: u    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
# s& _- u) _) r) X) ~. J9 C
# 开始测试
' O# {/ B$ b9 N- eevaluate(model_ft)

7 结语

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

1 h8 O  Q/ D6 [4 E0 }Epoch 0/2
----------
$ i  L# }/ i+ H$ V+ ^' R  utrain Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456
3 O& I! l1 M- G8 W
9 _5 x8 \8 Q) g, B; yEpoch 1/2
----------
$ L1 Z6 Z* |7 @- \* ~3 P& ltrain Loss: 1.1288 Acc: 0.6535
( u: m% Z+ Y! t' @* L% [$ a7 M8 }val Loss: 0.9012 Acc: 0.6929

$ |  {) g) L$ J1 b- N. O0 tEpoch 2/2
----------
4 w8 i$ ?- S4 a& R; wtrain Loss: 0.9578 Acc: 0.6706
5 ?5 \! z* t* H* j% Tval Loss: 0.8088 Acc: 0.6948
5 q% {& S, X* X  j2 }+ v3 f. I
Training complete in 6m 37s



- k0 r& D' }9 q/ |6 R
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

对于这个模型用 ResNet-50 作特征提取器会有更好的效果，不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
* n6 g/ j1 l  I% O5 W- Q
8 P! ^5 M; y4 C; k  v2 Q4 T- b5 |
( G5 v9 J: b+ T. I7 U! I语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

Deeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
& J  u9 V2 w6 {7 D( X. o* FUperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
U-Net 用于生物医学图像分割的卷积网络
: t8 k/ D1 t2 Q" y) \SegNet 用于图像分段的深度卷积编码器-解码器架构。
! n/ ~7 G' D, [: Q还有（DUC，HDC）、PSPNet等。

常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
! R% \- O  U6 L
对于损失函数，除了交叉熵误差，也可以用这些：

. o" J2 w" x: y: P8 YDice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
' s5 Q+ i. d# {5 A% G+ qFocal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
  Q+ Y$ B" j" u) r% iLovasz Softmax 查看论文：Lovasz - softmax损失。
/ ^- {$ V2 X: l# k7 O- W

# ^4 H4 n. ~9 h7 n9 g: S


0 F2 w+ @) K) U0 N7 _6 j" G————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
) E! S. C1 h1 w8 N) C! y# l5 m原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
: J" p2 B* ?7 F) L
. S% N% g  M7 n
& G' \1 J/ T9 r6 r0 f/ {0 }