Pytorch实战语义分割（VOC2012）

杨利霞

8 V' Y9 L  S+ i" \$ J+ W4 t
Pytorch实战语义分割（VOC2012）
2 f* @" U  V" m' T6 O: ~本文参照了《动手深度学习》的9.9、9.10章节，原书使用的是 mxnet 框架，本文改成了pytorch代码。
/ u. ^8 F8 p% j/ X1 x! t语义分割（semantic segmentation）问题，它关注如何将图像分割成属于不同语义类别的区域。值得一提的是，这些语义区域的标注和预测都是像素级的。
7 Y+ k" z5 d  z& o0 l

! K4 \  f. {6 K9 w. u语义分割中图像有关狗、猫和背景的标签
9 z' M. X1 Y7 T, Q; U# p文章目录
8 U( y+ S. D& U; z3 G( Q- @
% l' l! z& B& L/ t4 N1 图像分割和实例分割
, z! S+ L- K8 R) S& |! p2 Pascal VOC2012语义分割数据集
" m' {: Z( s& i# M, T, P+ R: ^2.1 导入模块
2.2 下载数据集
2.3 可视化数据
2.4 预处理数据
3 自定义数据集类
3.1 数据集类
- Z3 ]; t3 Q3 b, {$ C" T) s3.2 读取数据集
' U( T1 `/ a" N. Q4 构造模型
, V6 G* y4 b3 n8 N5 U4.1 预训练模型
4.2 修改成FCN
7 p5 v: p5 D" X" @% P4.3 初始化转置卷积层
: l9 Z' }& @, T8 p* Y! J) s5 训练模型
6 测试模型
6.1 通用型
6.2 不通用
7 结语
7 D; F* D/ u/ g, d% e1 y2 o1 图像分割和实例分割
9 p/ I  v6 R* [& D& c8 `
3 k) L2 Q3 h5 H$ k: \! ]2 G计算机视觉领域还有2个与语义分割相似的重要问题，即图像分割（image segmentation）和实例分割（instance segmentation）：
5 N- ]; v7 }3 w
图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息，在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入，图像分割可能将狗分割成两个区域：一个覆盖以黑色为主的嘴巴和眼睛，而另一个覆盖以黄色为主的其余部分身体。
9 `4 u% ~+ e4 p6 g8 C! q实例分割又叫同时检测并分割（simultaneous detection and segmentation）。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同，实例分割不仅需要区分语义，还要区分不同的目标实例。如果图像中有两只狗，实例分割需要区分像素属于这两只狗中的哪一只。
) G( C7 {5 t4 X  A8 A) Z
' ^* _/ R7 R" D7 m% v2 Pascal VOC2012语义分割数据集
1 H/ O& n* W. J
0 D% ^* v5 j( t' m. s, k3 ]- M& b2.1 导入模块
1 d6 u! [0 m( T; Z6 Kimport time
import copy
; F/ C5 `( h6 ?& Yimport torch
- B. d. a: |- B' d1 X3 zfrom torch import optim, nn
import torch.nn.functional as F
import torchvision
7 J% x* ^5 s9 Rfrom torchvision import transforms
, _+ Y( P& o+ [  k# _/ h/ X/ mfrom torchvision.models import resnet18
5 z+ z/ g& m7 B, d- ]0 u" himport numpy as np
" G* y, C' j# ]; rfrom matplotlib import pyplot as plt
from PIL import Image
import sys
sys.path.append("..")
from IPython import display
from tqdm import tqdm
import warnings
. K2 t3 M0 ?' W* ^warnings.filterwarnings("ignore")

2.2 下载数据集
& K& E/ d7 a: F1 s* k. L/ |
9 W4 y# t7 [4 O6 J" `1 I语义分割的一个重要数据集叫作Pascal VOC2012，点击下载这个数据集的压缩包，大小是2 GB左右，所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹，然后将其放置在data文件夹下，VOC2012文件目录是这样的：
- B; w4 Q" J, ]5 d- V1 o3 x* ?4 U

ImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
/ s. Z3 C) Z$ L1 wJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。
- y4 E( ~% G# L& b- I: D" b; T% O2.3 可视化数据

定义read_voc_images函数将输入图像和标签读进内存。
' N* U' s# f* Y9 m. R
def read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
+ H: V+ C9 A* E* ~    txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
2 q6 E  f1 W' f# p' ~    with open(txt_fname, 'r') as f:
5 s; ]+ ~% s2 `' y        images = f.read().split() # 拆分成一个个名字组成list
5 R& }) B% D. C& u5 I    if max_num is not None:
( z# H% I2 r) l1 y  E        images = images[:min(max_num, len(images))]
* o/ D) [# q+ i: v6 u) Q    features, labels = [None] * len(images), [None] * len(images)
/ ?! g2 e  ~7 n  {    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

- d# X: A& M0 s, g( u2 q4 \) q定义可视化数据集的函数show_images
4 D; }, ^: x0 P* {% F
9 w$ l# ]( R3 L0 m/ f+ x# 这个函数可以不需要
- p1 K5 s) W. P% s' rdef set_figsize(figsize=(3.5, 2.5)):
, L8 f+ t+ `, a* Z7 I7 ~    """在jupyter使用svg显示"""
    display.set_matplotlib_formats('svg')
) Y+ W; s$ \0 C/ [    # 设置图的尺寸
% v" c4 ^% H& u% ~" n+ a    plt.rcParams['figure.figsize'] = figsize
  ~+ E$ Y) C5 Y
" A: l  a5 e7 ]  S1 r! Ndef show_images(imgs, num_rows, num_cols, scale=2):
    # a_img = np.asarray(imgs)
    figsize = (num_cols * scale, num_rows * scale)
4 \( @5 D' M3 V/ _9 ?3 @& U9 N    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    for i in range(num_rows):
3 |! g2 D6 w& t7 E        for j in range(num_cols):
            axes[j].imshow(imgs[i * num_cols + j])
6 L* g6 {4 s* r, c            axes[j].axes.get_xaxis().set_visible(False)
6 s0 f- r, W. k1 v. {            axes[j].axes.get_yaxis().set_visible(False)
6 `# Q$ `6 R; v, W9 H( v    plt.show()
    return axes
5 m0 u. s7 h! S8 {- L
定义可视化数据集的函数show_images。

# 这个函数可以不需要
' L0 R! m# m/ y; n  }def set_figsize(figsize=(3.5, 2.5)):
. G& `; V9 Y8 v5 W9 }- A    """在jupyter使用svg显示"""
. H! t- M7 ?, _; O5 R! ~! J    display.set_matplotlib_formats('svg')
    # 设置图的尺寸
2 }9 c0 x" h( \9 _6 R" u    plt.rcParams['figure.figsize'] = figsize
: B; |8 M, J! k; L( o# T) i$ s2 H
) |& X) }$ W8 Udef show_images(imgs, num_rows, num_cols, scale=2):
0 L) o( a. n# G    # 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):
/ }# Z3 ]2 E( t1 ^        for j in range(num_cols):
7 @: O5 R9 J" ]! Y            axes[j].imshow(imgs[i * num_cols + j])
. v$ O; T: r4 _; `1 s+ S* S            axes[j].axes.get_xaxis().set_visible(False)
5 q. k3 Q9 R- [: ^# ~) D            axes[j].axes.get_yaxis().set_visible(False)
, y& d9 C& s" m    plt.show()
    return axes
画出前5张输入图像和它们的标签。在标签图像中，白色和黑色分别代表边框和背景，而其他不同的颜色则对应不同的类别。
+ x2 z+ H! l3 h' d1 ?
- y. g1 }$ c; ~( R0 G+ O# 根据自己存放数据集的路径修改voc_dir
( F2 O/ r1 x9 W. ?5 u, Dvoc_dir = r"[local]\VOCdevkit\VOC2012"
train_features, train_labels = read_voc_images(voc_dir, max_num=10)
n = 5 # 展示几张图像
; U/ `8 _7 C6 a- u; s; e- m* T* simgs = train_features[0:n] + train_labels[0:n] # PIL image
show_images(imgs, 2, n)
4 o! z$ [, g- w; r. ]

4 B0 p  e" v7 k5 P* b
列出标签中每个RGB颜色的值及其标注的类别。
, E0 |( B; [  `" x: R3 y, d# 标签中每个RGB颜色的值
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
2 f, a7 q' o7 v: [                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
. w! B2 h5 Z* C                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
1 u' x( @# a! Q" _  ?                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
! w1 B- s% p$ N6 {                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
# 标签其标注的类别
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
. u/ I9 F& c( N8 V/ M' ]               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
5 I+ [; R7 _2 f! m; x1 U7 w               'diningtable', 'dog', 'horse', 'motorbike', 'person',
8 L7 G; L) J) E               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
) @( q, I" H  t有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
2 _3 b5 u# ?: c- d- ?# D0 W% A
2 S2 y* B4 {3 E# A- p. a; B* N有了上面定义的两个常量以后，我们可以很容易地查找标签中每个像素的类别索引，voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来，当作 label 。
8 q- @: K4 M( r1 P) |3 O5 pcolormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
4 h  V' ^8 {" h! kfor i, colormap in enumerate(VOC_COLORMAP):
: S9 m/ I/ u- Z. |& j    # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
    colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
+ F( }# P2 g9 q8 r
! n: I  B7 S, W' j6 M, C" F# 构造标签矩阵
. M; G7 v9 N3 d+ e& v# @+ edef voc_label_indices(colormap, colormap2label):
    colormap = np.array(colormap.convert("RGB")).astype('int32')
* h9 ^$ `+ }7 d9 E    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
& a) V: ?  @+ S6 u
2 q( q$ v* K% d" E/ E; g可以打印一下结果
& Y+ f4 U* \, L& i
y = voc_label_indices(train_labels[0], colormap2label)
print(y[100:110, 130:140]) #打印结果是一个int型tensor，tensor中的每个元素i表示该像素的类别是VOC_CLASSES

2.4 预处理数据
/ g0 S2 e( T: z; v9 R! c
1 P& q2 I1 l% e, x6 f3 I3 P; f在语义分割里，如果使用缩放图像使其符合模型的输入形状的话，需要将预测的像素类别重新映射回原始尺寸的输入图像，这样的映射难以做到精确，尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说，我们使用图像增广里的随机裁剪，并对输入图像和标签裁剪相同区域。
: w, q' Y. G% S! M% {, ?

% r; Z2 @" o9 E5 U% W1 I3 ^" Q7 jdef voc_rand_crop(feature, label, height, width):
; }' s  i5 N2 f7 ~/ G    """
' m9 d- d. F5 n2 P    随机裁剪feature(PIL image) 和 label(PIL image).
1 W- K+ \& j7 \0 ?    为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
. q; C: J9 Z  v5 \# q  v    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.
+ L  k3 C; S/ ?5 q) t8 m0 N- r    Returns:
        tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
: k7 }; z2 D) v    """
; p4 c/ {1 z4 O& k    i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
/ O- L) o$ R' P2 i$ _: ?9 z* X8 A7 Z    feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
    label = torchvision.transforms.functional.crop(label, i, j, h, w)
! M2 M( p" l2 d. \! C1 q  w; {; @    return feature, label
# ?$ l$ e4 R' |" ~; Z, Q
# 显示n张随机裁剪的图像和标签，前面的n是5
; e/ Z$ q3 U" _3 R' |( T) E" w& ]" bimgs = []
for _ in range(n):
    imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
) T9 o/ n4 `' E, }; R! r; v- rshow_images(imgs[::2] + imgs[1::2], 2, n);

1 `6 f; m1 {3 C9 U# L' r8 ?2 W
2 G+ [) M0 x( t" C+ t& {" V) f: ~
0 J, y5 ^3 {' U


3 自定义数据集类
* t  h, l# g' y+ W! o
3.1 数据集类
# A( P2 p% h9 v. t3 A$ F  L3 K" R
torch.utils.data.Dataset是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法
! Z( `2 H! [% Y& A+ a6 ~0 A
' N( `7 i! O5 b__len__ 实现 len(dataset) 返还数据集的尺寸。
__getitem__用来获取一些索引数据，例如 dataset[idx] 中的(idx)。
2 u& j: ^& I7 g! U2 d由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。此外，因为之后会用到预训练模型来做特征提取器，所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
3 Z; V! J; \( k' F& [
  E7 w, k$ h; v% v' \6 f8 mclass VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):
+ a# r5 i) }! A: h7 z3 ~        """
! @/ Q' e3 A# m( w5 D        crop_size: (h, w)
        """
        # 对输入图像的RGB三个通道的值分别做标准化
8 ~/ @' G) e1 p) k% F* M        self.rgb_mean = np.array([0.485, 0.456, 0.406])
) z1 P0 _* d6 X6 a        self.rgb_std = np.array([0.229, 0.224, 0.225])
1 v$ m1 l1 C6 _; g. @        self.tsf = torchvision.transforms.Compose([
: B$ a$ p3 K7 \            torchvision.transforms.ToTensor(),
7 p. Y& e! q- a            torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])
1 n* Y6 u0 O; K5 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函数所移除
7 [( B! E8 m! e+ p: r        self.features = self.filter(features) # PIL image
        self.labels = self.filter(labels)     # PIL image
; [" m( ~' Q* j5 ]5 n2 R        self.colormap2label = colormap2label
+ R; r; l5 U3 [' D4 J$ L! H6 t        print('read ' + str(len(self.features)) + ' valid examples')

! N! ^8 h+ u7 K8 `4 w    def filter(self, imgs):
        return [img for img in imgs if (
$ Y' r1 H. |5 Q3 h' `- u+ M            img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]

. Q* W1 i2 A1 y# K9 s  k    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)
5 a' o1 i' v7 ^7 R                                # float32 tensor           uint8 tensor (b,h,w)
        return (self.tsf(feature), voc_label_indices(label, self.colormap2label))
% z) C8 _( b- M4 s+ y
" R  [2 k. {( j# B0 n    def __len__(self):
        return len(self.features)
3.2 读取数据集

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

$ o0 U8 y; F* _0 d% k! Obatch_size = 32 # 实际上我的小笔记本不允许我这么做！哭了（大家根据自己电脑内存改吧）
" Y; h- q  D$ h' F  Q4 P; mcrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)
' Y% C+ {2 [; z$ i# \max_num = 20000 # 最多从本地读多少张图片，我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~
3 B3 h" ^( C  ~8 ?: |( v
# 创建训练集和测试集的实例
voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
9 V! _+ m; u) t1 S- P0 P6 Z
# 设批量大小为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)
% ]7 V- O$ N3 k1 ]- ^" Ztest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
! q& o9 {: [5 O9 |; Y3 }                             num_workers=num_workers)

4 g1 |+ |$ y9 X/ A% E4 h6 N# 方便封装，把训练集和验证集保存在dict里
dataloaders = {'train':train_iter, 'val':test_iter}
3 N+ I. ?6 d1 Adataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}
. y, D) ~& F6 \. O+ ~0 c3 m5 ?$ @& Z
4 构造模型4.1 预训练模型

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
5 p" K! a7 \4 c1 X: w
( @- h3 ]" e, c# v) d# tnum_classes = 21 # 21分类，1个背景，20个物体
2 N3 F) Q, o$ l9 g. j% H, dmodel_ft = resnet18(pretrained=True) # 设置True，表明要加载使用训练好的参数
8 f& N0 S9 r2 b" C. }, o
# 特征提取器
: ^: H8 |# R. i, afor param in model_ft.parameters():
    param.requires_grad = False
4 S# o& X, ^9 c/ |6 A4.2 修改成FCN
( N2 C9 l5 j; b/ [
全卷积⽹络（顾名思义全部都是卷积层）先使⽤卷积神经⽹络抽取图像特征，然后通过 1×1​ 1\times 1​1×1​ 卷积层将通道数变换为类别个数，最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同，并在空间位置上⼀⼀对应：
1 Q$ G# y, _2 S  l* t6 a/ w最终输出的通道包含了该空间位置像素的类别预测。
, R6 P9 l- R" I$ h) N* q9 k- x
0 N% h# B7 x% V6 j9 }对于转置卷积层，如果步幅为 S​ S​S​、填充为 S/2​ S/2​S/2​ （假设为整数）、卷积核的⾼和宽为 2S​ 2S​2S​，转置卷积核将输⼊的⾼和宽分别放⼤ S​ S​S​ 倍。
' {% q) v+ E9 f% D- A/ N
可以先打印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 倍即可。

& V, A1 @, O$ \8 c8 k+ smodel_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层
4 n1 y# l. Q0 _8 m. U0 \3 ~) M              nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class
4 V, d* q1 w4 _, _+ R) T5 x; B/ ^, ^              nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小
3 @: F- x) ], y! h4 i# ?' i6 c0 f
# 对model_ft做一个测试
x = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据
$ O: J& t& w& u# l) m5 Fprint(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])

# 打印第一个小批量的类型和形状。不同于图像分类和目标识别，这里的标签是一个三维数组
# for X, Y in train_iter:
& K/ {6 b( P: _' k9 {#     print(X.dtype, X.shape)
#     print(Y.dtype, Y.shape)
/ B6 ^  s% `, U) X+ f  ^#     break


1 s" L$ S0 f$ p" X$ n1 U5 h4.3 初始化转置卷积层

在图像处理中，我们有时需要将图像放⼤，即上采样（upsample）。上采样的⽅法有很多，常⽤的有双线性插值。简单来说，为了得到输出图像
. w5 }1 P2 d+ Q2 K  `4 l# }. A在坐标 (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函数构造的卷积核的转置卷积层来实现。
' K/ Y4 C$ S4 P$ ~+ v5 d
' v2 Q! Q# m+ ^! o$ l9 r0 {
# 双线性插值的上采样，用来初始化转置卷积层的卷积核
def bilinear_kernel(in_channels, out_channels, kernel_size):
3 z9 y- o7 p2 U* H1 N$ \; R2 j0 X    factor = (kernel_size+1)//2
; g% c) b* `+ C5 ]" W$ E    if kernel_size%2 == 1:
        center = factor-1
    else:
$ \# B0 l0 R5 p; L        center = factor-0.5
; ]+ v7 @% H  d2 L8 s, ]    og = np.ogrid[:kernel_size, :kernel_size]
    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')
9 @; r5 {/ D* G  H    weight[range(in_channels), range(out_channels), :, :] = filt
    weight = torch.Tensor(weight)
4 g+ {! J9 U6 @/ T: u% g2 I    weight.requires_grad = True
- L$ d2 S) w, [4 S    return weight


在全卷积⽹络中，将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层，采⽤Xavier XavierXavier随机初始化。

2 u5 q  W# l. k8 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)

) u0 \; E/ Q' G% L
1 J( q, F, i, W# M
5 训练模型

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

0 M5 B* D% B9 r4 idef train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
0 P2 r' z/ g% l* P$ `6 W    best_acc = 0.0
) v" u/ I5 f, }$ |& l8 C2 J" i    # 每个epoch都有一个训练和验证阶段
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
, x9 L1 d5 G3 p4 ]; i; P$ b( m% W$ {        for phase in ['train', 'val']:
            if phase == 'train':
/ }8 O" m9 ]. J. Z7 U3 ]                scheduler.step()
! m/ P& T# P7 L) ?8 j                model.train()
) e+ f2 G& \) g( U( v2 o" c            else:
                model.eval()
1 f8 U3 x6 P, x; `5 n4 V5 c9 c4 e( F            runing_loss = 0.0
            runing_corrects = 0.0
: y3 O- G, s6 C/ r: b% d" n( l            # 迭代一个epoch
% g6 R2 y7 n- D8 p4 b5 z            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(device), labels.to(device)
4 u/ ~7 `6 Q* U% ]                optimizer.zero_grad() # 零参数梯度
                                # 前向，只在训练时跟踪参数
& P" X8 Z  M8 v6 R& _1 e                with torch.set_grad_enabled(phase=='train'):
: S  |$ b$ r7 y3 F0 ~' U4 \9 C0 A                    logits = model(inputs)  # [5, 21, 320, 480]
# |9 O) k2 T% A7 r7 E                    loss = criteon(logits, labels.long())
                    # 后向，只在训练阶段进行优化
9 q3 I: h: ]! B( m                    if phase=='train':
. M$ e  U" H' }7 M  j( t) [% w                        loss.backward()
6 y. M: S8 T5 ]( [5 P                        optimizer.step()
                                # 统计loss和correct
                runing_loss += loss.item()*inputs.size(0)
  V# @$ T2 W+ i& d% w                runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)

' ]3 O5 Z6 n, j% [) V            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))
" d. v* _# E9 t. P* c                        # 深度复制model参数
6 M! m" H4 E* B+ ~' Y* q' W6 _  p            if phase=='val' and epoch_acc>best_acc:
+ P3 b) e; ~( K9 ]                best_acc = epoch_acc
' u/ Q1 O* c7 y                best_model_wts = copy.deepcopy(model.state_dict())
        print()
# M0 x, |! Q; ?* \% X( ~    time_elapsed = time.time() - since;
. N. C6 C5 @" M/ O8 Z; ~2 r    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
/ N, s- I$ L) M$ U( z0 W6 ^& N    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
; ]4 F# M: I/ _    return model

下面定义train_model要用到的参数，开始训练
0 J. _9 w6 d& u- ^2 b
% `3 X. Q( [( ?: l) Kepochs = 5 # 训练5个epoch
criteon = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
2 z- a7 M5 B) |/ c/ i# 每3个epochs衰减LR通过设置gamma=0.1
2 a5 b! f, [% g9 H; Texp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
1 r7 L' u- ~2 s7 _8 P6 J7 N
) i! N* ~, q. N2 \/ N. o% ^# 开始训练
model_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)

6 测试模型

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

def label2image(pred):
, M9 s6 C9 c0 r4 s* z. M5 }    # pred: [320,480]
6 S$ v$ x. z5 |  ]    colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
, C# |8 Z( w# F$ L4 ]& q    x = pred.long()
    return (colormap[x,:]).data.cpu().numpy()
# e/ b' c* G- R* m# h9 ^0 `/ \2 j

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

6.1 通用型

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

mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
% o$ I9 f4 D( A7 `( r1 dstd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)
def visualize_model(model:nn.Module, num_images=4):
: b* X7 i9 F1 p6 M    was_training = model.training
; w/ \: K; D; y, f3 z4 f    model.eval()
    images_so_far = 0
    n, imgs = num_images, []
    with torch.no_grad():
/ A7 E  _- P" B* \        for i, (inputs, labels) in enumerate(dataloaders['val']):
6 F6 a0 f7 g* T* Y            inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
            outputs = model(inputs)
4 M5 U8 y* _2 K9 ?            pred = torch.argmax(outputs, dim=1) # [b,320,480]
! G1 h( ]- v, p# _            inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
6 ], I2 X! f+ x. A) R
            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)
6 [( b& J# D$ @$ y$ b                    # 我已经固定了每次只显示4张图了，大家可以自己修改
3 z' t1 E9 g* P2 E                    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
                    return model.train(mode=was_training)
1 b0 V& S1 p3 A9 a2 i  v# `0 D
# 开始验证
' v( L1 V- T$ d, gvisualize_model(model_ft)
6.2 不通用

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

# 预测前将图像标准化，并转换成(b,c,h,w)的tensor
6 r" M: O, G) @  ]* D/ E* ?* adef predict(img, model):
- i, o" v2 ^" Z: z+ y7 k: \    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)
    pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)
; E" T% G0 t2 {8 K! h) \    return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
8 C& f: O9 v4 y9 t  l; {
, t6 p& N* M4 a0 Y: b3 i3 gdef evaluate(model:nn.Module):
) F* t3 u( Z) P$ A* S: y: Z    model.eval()
$ V  a* ~% L, C7 g! P0 L4 D( B$ Q    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
; w& C! R6 f; s3 ~- I        pred = label2image(predict(xi, model))
$ m( [% m5 i/ o6 z; w        imgs += [xi, pred, yi]
    show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
* D/ d- B4 f2 B, ^; Y8 d" C5 a6 p
) L6 ?  x* ?! d$ R# 开始测试
evaluate(model_ft)
9 @' ^6 b; g' o3 j2 |
2 p# t; U; p0 c9 ?# e7 结语

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

Epoch 0/2
  `4 u+ D8 ?4 Z5 g5 s8 l----------
/ t# H2 n0 s# d5 F! r2 h2 _0 ytrain Loss: 1.7844 Acc: 0.5835
val Loss: 1.1669 Acc: 0.6456
5 `5 {8 h, t2 R# [& Q8 D( I
4 d, V6 |/ W' S' X0 Q  d# K7 A! BEpoch 1/2
  I( y; z: Z7 L6 g3 j0 N9 m" `! ]----------
) O- @7 a9 ]9 a2 R0 xtrain Loss: 1.1288 Acc: 0.6535
8 w5 ?. K3 K: S% ?- Gval Loss: 0.9012 Acc: 0.6929

Epoch 2/2
7 G/ e, Q" D( B0 X, q: o, n% ^----------
train Loss: 0.9578 Acc: 0.6706
val Loss: 0.8088 Acc: 0.6948
; z: K$ N( n- G6 I/ m4 @
8 K9 C' c, A1 F5 ~; DTraining complete in 6m 37s
" p7 M! }( H: H; k0 v0 b

9 f( y9 t! }) ]# z
9 |3 C6 {# O2 U, T) t
当 epochs = 5 时，训练集的精度在 89 8989% 左右，测试集的精度可以达到 86​ 86​86​ %。

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


语义分割还有很多可用的模型，本文用的是 FCN，在其它一些模型上会有更好的表现：

/ Z& X  w) y- A9 b, ?1 C3 B/ ^8 XDeeplab V3+ 具有可分离卷积的编码器/解码器，用于语义图像分割[论文]
GCN 通过全局卷积网络改进语义分割[论文]
UperNet 统一感知解析
ENet 用于实时语义分割的深度神经网络体系结构[论文]
) x; B4 V1 m6 _2 cU-Net 用于生物医学图像分割的卷积网络
SegNet 用于图像分段的深度卷积编码器-解码器架构。
# x( G/ c8 d$ ~7 Y/ R/ Z; b+ O) X& V还有（DUC，HDC）、PSPNet等。

常用的语义分割数据集也有很多：Pascal VOC、CityScapes、ADE20K、COCO Stuff等。
& ?: L! ^% l# [: u. \3 R
' G  _9 M$ |; ^# X: Q对于损失函数，除了交叉熵误差，也可以用这些：
- N. `3 G. X$ K& w
Dice-Loss 可以测试两个样本之间的重叠度量，可以更好地反映训练目标，但该损失函数具有很强的非凸性，很难优化。
CE Dice loss Dice 损失与 CE 的总和，CE 提供了平滑的优化，而 Dice 损失则很好地表明了分割结果的质量。
Focal Loss CE 的另一种版本，用于避免类别不平衡而降低了置信度的情况。
" \8 O6 _( h2 c1 j8 ?; ^- d$ }Lovasz Softmax 查看论文：Lovasz - softmax损失。

2 r* u: c, U, o  b3 C3 B8 p7 j



  k6 b$ L2 H4 g————————————————
版权声明：本文为CSDN博主「小红不吃糖」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_43280818/article/details/105916507
9 i6 Y# O4 K" i/ i" D
  ^2 [2 S$ x5 F' f: m4 ^) o