5 h% W) G3 |: J4 G. m# P6 X
8 Q5 c. ^& s+ x$ NPytorch实战语义分割(VOC2012)* ~( A# S! R- N
本文参照了《动手深度学习》的9.9、9.10章节,原书使用的是 mxnet 框架,本文改成了pytorch代码。
I% a% ~' ?3 U. n5 ~语义分割(semantic segmentation)问题,它关注如何将图像分割成属于不同语义类别的区域。值得一提的是,这些语义区域的标注和预测都是像素级的。; f$ p/ J0 z4 L1 W) z9 F: b
' D* y1 s+ f1 f
5 f+ h5 k3 q- S5 X语义分割中图像有关狗、猫和背景的标签
& r8 \; n; I& t9 r1 H文章目录
3 u. R1 q6 h* e. E% n" u' Y+ B
1 图像分割和实例分割
. E" ^* J* r' X2 G6 Y- m2 Pascal VOC2012语义分割数据集) o" g' ?! H" P
2.1 导入模块
- f9 S1 n2 b( b* B9 W2.2 下载数据集
3 C& V6 M' J2 L2.3 可视化数据
$ U* f8 A4 I* h9 _$ S2.4 预处理数据
; y7 u0 p; X: \$ r; W7 B3 自定义数据集类
. p6 A' N% l) V* F6 z3.1 数据集类: e; {3 I: j: k. c, u
3.2 读取数据集
; l. ~% r2 z$ T5 ~5 K0 [4 构造模型. T0 K& k7 H: c& m* j# h
4.1 预训练模型8 F' b/ y9 y. t3 j( C
4.2 修改成FCN& C; g0 W6 K9 @4 U( x& F% N8 B
4.3 初始化转置卷积层( S- W" D9 i: f$ j) G5 A8 L6 V; e
5 训练模型) }) r9 a/ {/ B4 B
6 测试模型( }# d* Z7 T% S7 r! s3 P
6.1 通用型
. J5 N0 {' { G6.2 不通用# N: J( C0 z, G( D- B1 E
7 结语6 g7 b/ T' Z* k$ ~+ ?
1 图像分割和实例分割' d' V& b# X; ~
3 }# c0 N# S$ R3 Y+ D5 p计算机视觉领域还有2个与语义分割相似的重要问题,即图像分割(image segmentation)和实例分割(instance segmentation):$ @" }2 s; m; N0 s. U8 c/ z( t1 H
# n5 i) ]6 V5 ]( U" V, R: n图像分割将图像分割成若干组成区域。这类问题的方法通常利用图像中像素之间的相关性。它在训练时不需要有关图像像素的标签信息,在预测时也无法保证分割出的区域具有我们希望得到的语义。以上图的图像为输入,图像分割可能将狗分割成两个区域:一个覆盖以黑色为主的嘴巴和眼睛,而另一个覆盖以黄色为主的其余部分身体。4 ?3 C! z4 L. J; T0 L
实例分割又叫同时检测并分割(simultaneous detection and segmentation)。它研究如何识别图像中各个目标实例的像素级区域。与语义分割有所不同,实例分割不仅需要区分语义,还要区分不同的目标实例。如果图像中有两只狗,实例分割需要区分像素属于这两只狗中的哪一只。
) X9 E t t6 B& L% T; h1 A- b/ Z- t z% U9 m5 j% x, k9 C( @( M! p
2 Pascal VOC2012语义分割数据集$ ~, s* C8 k2 W8 [5 l& Q
8 k# }* v, N9 S5 ~5 D2.1 导入模块, B$ t% k( b3 d3 @7 K1 V9 }
import time
! V+ f6 x z: a4 `' oimport copy
9 y- ]$ P( m W4 Z; }: ximport torch
0 j" g6 V6 t! ~! v" S1 L0 h: _6 ^7 cfrom torch import optim, nn
, g8 s3 @0 F* T, A. W! E; Qimport torch.nn.functional as F
2 T0 F( D* W# G9 kimport torchvision( Z% S' F% u7 q. A5 k
from torchvision import transforms+ K( N* H& i5 Y1 E5 f8 Y( ^
from torchvision.models import resnet180 l# O% K& r, E; K
import numpy as np9 j5 y6 N4 K0 w" D- j6 m! p" e
from matplotlib import pyplot as plt
: n# ~5 I' R" Y. A6 P4 Dfrom PIL import Image
3 F: e2 m |5 oimport sys
1 l$ n6 }) K$ Q# d% Q; G' [+ v; Isys.path.append("..")
5 r- v, u# L, D$ { m9 ^ a8 zfrom IPython import display6 E- v' I7 Q9 Q2 |/ ?
from tqdm import tqdm# d% A- x+ o) Z& j' F% D" \' T2 l
import warnings( j' i$ D/ n; ^5 V" l
warnings.filterwarnings("ignore")8 z3 W `! }. @% ?6 U, c0 T2 z# r7 _0 r
2 y+ R: x( O% _2.2 下载数据集8 q0 V- ? z9 Y/ I8 c
) C, y9 w& u8 T5 v; D9 o, M G2 E
语义分割的一个重要数据集叫作Pascal VOC2012,点击下载这个数据集的压缩包,大小是2 GB左右,所以下载需要一定时间。下载后解压得到VOCdevkit/VOC2012文件夹,然后将其放置在data文件夹下,VOC2012文件目录是这样的:
3 s2 g& z2 \$ U3 P+ ~, f8 Z( Q, F X; a1 q
$ G8 E" y4 W+ o$ C# rImageSets/Segmentation路径包含了指定训练和测试样本的文本文件
" S! H4 p( } LJPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式,其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。0 T1 z( ^+ j6 m9 k1 J. [' J
2.3 可视化数据4 O/ r Z, b4 o/ t' v" R: f/ ]
. S+ l/ P( H( P$ g1 q
定义read_voc_images函数将输入图像和标签读进内存。' D1 j" f% F5 K- a: ?9 _: C' W
3 I+ W, b3 F2 cdef read_voc_images(root="../../data/VOCdevkit/VOC2012", is_train=True, max_num=None):
: \( ~, W; o4 E9 }$ @1 D* i7 L% V' x txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt')
# g( D6 j1 @5 S9 w with open(txt_fname, 'r') as f:
3 ~4 B8 j8 R8 e$ }2 _9 r @# ? images = f.read().split() # 拆分成一个个名字组成list
# M& O0 x2 g* v N. u if max_num is not None:" q$ t* B" ]. }
images = images[:min(max_num, len(images))] r; i9 v4 O! C* _
features, labels = [None] * len(images), [None] * len(images)
9 [( j$ w$ e, M! T3 D- k& O! ^ for i, fname in tqdm(enumerate(images)):: O N! `3 C% a7 c* l' a, ^
# 读入数据并且转为RGB的 PIL image
1 C9 b7 t- F4 r7 X% s features = Image.open('%s/JPEGImages/%s.jpg' % (root, fname)).convert("RGB")
1 R' m y( g3 J, m! b labels = Image.open('%s/SegmentationClass/%s.png' % (root, fname)).convert("RGB")
8 F% f" c+ A+ ]: \+ i return features, labels # PIL image 0-255
# R* j" ~9 C9 ^3 B% r% ^) f8 _2 ?& @
定义可视化数据集的函数show_images8 {( p1 K+ I1 f
2 |) O! t. a. s4 W: e# 这个函数可以不需要
& Y' p. d+ e6 A2 tdef set_figsize(figsize=(3.5, 2.5)):
0 }2 p. G7 ~# n" u( o1 {9 s) W """在jupyter使用svg显示"""
W4 a* `0 g) p. c9 u- y display.set_matplotlib_formats('svg')5 f$ y( u& n' o+ X
# 设置图的尺寸
& [- P7 C- F% I8 y( l% K plt.rcParams['figure.figsize'] = figsize2 m* U6 x+ w1 V c9 P' ]: `
. D' ^( w7 _: a- Tdef show_images(imgs, num_rows, num_cols, scale=2):
# N; c$ Y/ d: M1 i # a_img = np.asarray(imgs)9 B9 Z i' O$ W8 @
figsize = (num_cols * scale, num_rows * scale)0 d$ T; K) e) e0 h) x4 X5 C, a
_, axes = plt.subplots(num_rows, num_cols, figsize=figsize)% B0 p. m0 d$ x% q; A8 t
for i in range(num_rows):2 {( M/ S# H7 D& k) J! e* V
for j in range(num_cols):+ [/ h. h8 E1 e1 ^
axes[j].imshow(imgs[i * num_cols + j])8 c1 x1 |) T% }+ ]7 b
axes[j].axes.get_xaxis().set_visible(False)
. y! s; ? e. H- f/ B4 Z. b0 R% X5 F axes[j].axes.get_yaxis().set_visible(False)
% E6 ~4 j8 N2 C6 r# R( R' e& T1 t plt.show()# s! U! u5 S5 _) Q7 t5 {6 S
return axes% k+ n' a3 K# f2 H
6 C1 T) i4 W3 _8 l& l8 \定义可视化数据集的函数show_images。
/ W" ~5 h& f1 D. n% U, Q7 N6 q
2 ?) C3 ]* Z( |& L& W1 ]8 t/ @! j# 这个函数可以不需要; r4 }' Q1 M- h& S) g4 o
def set_figsize(figsize=(3.5, 2.5)):7 h" o) n9 c$ r8 Y o7 I+ z
"""在jupyter使用svg显示"""! Q2 t2 A! E$ x5 ~- ^$ n& w: P9 G
display.set_matplotlib_formats('svg')# f) |; ~) |: \( m- h
# 设置图的尺寸
, ^ T9 Z4 A q plt.rcParams['figure.figsize'] = figsize
, s! K) X" L5 v# f
2 U8 v5 }7 g# e6 o/ L' l7 ]) @def show_images(imgs, num_rows, num_cols, scale=2):
! A2 d! d# c! v6 X2 v% |7 T1 E # a_img = np.asarray(imgs): n K( q. H7 D4 ~$ j- D7 x
figsize = (num_cols * scale, num_rows * scale)
; ~8 l& E7 r" I" i0 s _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
9 z+ {. J' h- n X% R for i in range(num_rows):
9 z w- T7 I. X4 B& \8 M( G for j in range(num_cols):
' F% n5 L2 X# r* v4 U+ ?' Y axes[j].imshow(imgs[i * num_cols + j])5 c6 c- z9 n$ Q, h5 r$ B# M
axes[j].axes.get_xaxis().set_visible(False)* B5 V" U t" f! i. E' F+ k
axes[j].axes.get_yaxis().set_visible(False)/ R; t. R# g, { M4 [
plt.show()
$ H$ Q* i: h/ s( p return axes) O& B& ~# K; X0 M6 M
画出前5张输入图像和它们的标签。在标签图像中,白色和黑色分别代表边框和背景,而其他不同的颜色则对应不同的类别。
8 J% h% \6 w6 I# r# T6 u
; ~+ h; N# g& N' A; ~0 I0 p- z5 H# 根据自己存放数据集的路径修改voc_dir% A. G! S" q9 h) U V. ^
voc_dir = r"[local]\VOCdevkit\VOC2012"
0 Z2 O4 {$ B( `. a* htrain_features, train_labels = read_voc_images(voc_dir, max_num=10)7 g% h, h+ G0 z0 b0 V3 m6 A- I# c! M
n = 5 # 展示几张图像+ D% b" ~, b* H3 D' [, @2 B% ^
imgs = train_features[0:n] + train_labels[0:n] # PIL image7 G; g5 H+ n1 z: k E; v- K# a& L8 z
show_images(imgs, 2, n)
+ N, A# v. S8 K- F+ u3 k4 q+ k* c( h1 q: d% P% F
: ?3 {8 O# [+ y# Z0 n
. ^$ O/ L- C% e列出标签中每个RGB颜色的值及其标注的类别。
( z; ~$ p7 P7 H# q1 m# 标签中每个RGB颜色的值* u" T H2 J) |$ N
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],5 ?6 w% V/ A. S7 I
[0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],3 [; u4 {! ?" s4 r' u% c ~
[64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],# T1 S3 `# J6 P/ e/ f, ^# g8 Z" k
[64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],. ?' m; ~$ s$ N) u9 |8 {
[0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],6 o' V9 t* ^8 A6 S, b! V. ~
[0, 64, 128]]: b) L, U) {7 H7 @+ M5 }) }
# 标签其标注的类别1 C2 ^$ z6 b D$ _
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',$ _. w% e& y. V3 p
'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
! V1 D. y3 l+ `1 l+ t- G9 }) j2 M/ |( l 'diningtable', 'dog', 'horse', 'motorbike', 'person',
8 u/ s- h, D9 Z W% Y1 I$ r1 y A 'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']: r! p3 n2 E, V x2 F
有了上面定义的两个常量以后,我们可以很容易地查找标签中每个像素的类别索引,voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来,当作 label 。
1 r* M+ x3 ], J' x Z( F* k2 m- K9 Q5 ]6 Q% U4 }: X6 T
有了上面定义的两个常量以后,我们可以很容易地查找标签中每个像素的类别索引,voc_label_indices是根据colormap2label把标签里的 rgb 颜色对应上面的VOC_COLORMAP中的下标给取出来,当作 label 。0 x7 d$ _& k. r
colormap2label = torch.zeros(256**3, dtype=torch.uint8) # torch.Size([16777216])
4 n3 _ ]6 ?! \for i, colormap in enumerate(VOC_COLORMAP):2 y) D2 t! o& l8 j3 O
# 每个通道的进制是256,这样可以保证每个 rgb 对应一个下标 i
- S. k& L4 ^' E. a1 b1 N- E colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i$ d) T4 R+ ~5 ?' U+ X
- \* W0 H, f0 {2 _$ m* |8 \! u# 构造标签矩阵
. c' I3 w1 I5 Udef voc_label_indices(colormap, colormap2label):
3 |0 c5 @$ m) \) { colormap = np.array(colormap.convert("RGB")).astype('int32')# ^( V1 w8 m1 ~# x- _ @
idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2]) " c8 H7 B: N3 i1 T; M' e" [& Z% f
return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标
( N0 i$ M0 y1 }+ g) a9 G9 N& F) x, Y6 z" {& R( o% Q
可以打印一下结果
, {4 M7 S" z. o$ m$ M# d8 l. w2 s3 p- _0 `$ b9 c
y = voc_label_indices(train_labels[0], colormap2label)9 ]: x6 a$ e' w, |' g
print(y[100:110, 130:140]) #打印结果是一个int型tensor,tensor中的每个元素i表示该像素的类别是VOC_CLASSES* `' U6 K; \" G1 n
7 H4 \7 a) h* O A% {
2.4 预处理数据# r7 e$ w( t! L6 n! \: r9 M
' v' i- ~ p! E8 F2 l( ~* o在语义分割里,如果使用缩放图像使其符合模型的输入形状的话,需要将预测的像素类别重新映射回原始尺寸的输入图像,这样的映射难以做到精确,尤其是在不同语义的分割区域。所以选择将图像裁剪成固定尺寸而不是缩放。具体来说,我们使用图像增广里的随机裁剪,并对输入图像和标签裁剪相同区域。, O- ], q2 @+ u9 E$ M9 \, h
" O: ]" J7 {0 Z. `5 A* ~& z* l J$ s& H
def voc_rand_crop(feature, label, height, width):( c* h! n$ x: Q
"""( X# E$ k3 T/ O4 T6 y! G( L8 U2 l
随机裁剪feature(PIL image) 和 label(PIL image).
3 [7 |$ H! t( y9 l 为了使裁剪的区域相同,不能直接使用RandomCrop,而要像下面这样做
7 ~3 Y( A) B! B, J8 D1 R7 [ Get parameters for ``crop`` for a random crop.' W( R5 U' L; ?. o
Args:7 { X% U4 U! p5 X% b6 U' ^* i- z
img (PIL Image): Image to be cropped.
# O+ o s: I: c2 L; }7 k0 ~ output_size (tuple): Expected output size of the crop.$ z3 C7 F9 q. p( I
Returns:' H* g$ W& v8 ?. R1 Q
tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
) E1 e# S6 d, v4 f# L """
* ^* B; @( r. @& K! l" b# a i,j,h,w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))# N6 v, P0 F$ V# ^1 G; N* m
feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
$ p5 d) S- q6 w- D* c: q, v8 R label = torchvision.transforms.functional.crop(label, i, j, h, w)+ f8 I9 l) |$ `# b% h( R0 _
return feature, label3 M$ A6 z, ^0 L, w) k, d) {, @' Y( P
! F, \* x7 `$ T6 q3 q, W# 显示n张随机裁剪的图像和标签,前面的n是5( E" V# F3 D6 @* j+ r ^
imgs = [] P; Q$ C1 a1 D1 v9 z5 H3 p
for _ in range(n):4 f2 |' C$ _: v; R" Q
imgs += voc_rand_crop(train_features[0], train_labels[0], 200, 300)
/ {, k" Y. w/ x5 E' F& Kshow_images(imgs[::2] + imgs[1::2], 2, n);
1 p7 h' `$ p: ]: c
% k V' W a" a- A l5 |, H* l% E. {
5 Q/ V' z9 [$ O+ K
) p; H2 p$ ]' @
! t* ~7 O$ t- r: u- V% |4 C9 k1 m% {# _7 X
3 自定义数据集类9 |* r# j5 f9 D, w- J* h
" [0 V0 P: q5 ^+ `6 B) B/ {3.1 数据集类
- I, k' i4 m. W3 q0 P. x& K1 C/ f8 m! V1 z4 q
torch.utils.data.Dataset是表示数据集的抽象类,因此自定义数据集应继承Dataset并覆盖以下方法( Q; b. j( r) ~" e% z
9 S; h) W5 K [, P__len__ 实现 len(dataset) 返还数据集的尺寸。& @0 [9 D, V4 b* ^
__getitem__用来获取一些索引数据,例如 dataset[idx] 中的(idx)。
, ^* {; E4 l, f6 }由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸,这些样本需要通过自定义的filter函数所移除。此外,因为之后会用到预训练模型来做特征提取器,所以我们还对输入图像的 RGB 三个通道的值分别做标准化。
, a; M2 P* C+ [$ L$ D) e* D6 l7 v! x. O' \& I
class VOCSegDataset(torch.utils.data.Dataset):8 b- ?$ h6 r3 N8 Q6 O# e
def __init__(self, is_train, crop_size, voc_dir, colormap2label, max_num=None):8 Z% l* }9 D! |9 o1 k6 u* G$ x
"""
" U* g4 b! _( F+ s- R* i' O% c crop_size: (h, w)
" v3 X3 G5 e) X [' X0 y2 R# b" } """
0 }) [- W1 U" A# C. ] # 对输入图像的RGB三个通道的值分别做标准化
' G( e/ o, s$ [ ^1 v% r$ l self.rgb_mean = np.array([0.485, 0.456, 0.406])
3 x6 p1 E8 o3 ] a( c, C+ }. U self.rgb_std = np.array([0.229, 0.224, 0.225])
# u) M9 G! I/ P+ i6 M self.tsf = torchvision.transforms.Compose([# r' \6 i6 E0 H/ Z/ i" P3 M
torchvision.transforms.ToTensor(),
* Q5 h# l- D: _; p2 m torchvision.transforms.Normalize(mean=self.rgb_mean, std=self.rgb_std)])! y' W8 L# y* I( Q% r
self.crop_size = crop_size # (h, w), h/ a! c" \1 p4 E
features, labels = read_voc_images(root=voc_dir, is_train=is_train, max_num=max_num)
; |! Q) c7 T8 t# 由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸,这些样本需要通过自定义的filter函数所移除! r1 W$ _: G4 X' C
self.features = self.filter(features) # PIL image; I8 `7 n9 p! H- [/ D7 m
self.labels = self.filter(labels) # PIL image
! S6 `9 l5 a8 q& Q9 X2 g. e, F self.colormap2label = colormap2label5 K" w e% M$ t
print('read ' + str(len(self.features)) + ' valid examples')
$ I! z' q$ o; J* r
3 V3 a, N& c; @ \" u" ]; w def filter(self, imgs):5 Q* X! B$ {, e* d$ d$ L1 }
return [img for img in imgs if (
) g, q( ]. y1 }& C img.size[1] >= self.crop_size[0] and img.size[0] >= self.crop_size[1])]
' M) ^1 R/ j. B* @- m3 e0 W; x
def __getitem__(self, idx):
6 s) P, [; C; F2 @( A# m' ` feature, label = voc_rand_crop(self.features[idx], self.labels[idx], *self.crop_size)+ o I$ `2 X7 X( I. @
# float32 tensor uint8 tensor (b,h,w)4 V2 Q" C. X9 |, x+ ?; |, W
return (self.tsf(feature), voc_label_indices(label, self.colormap2label))5 F8 y! C( |7 s+ R( _9 y; I: q5 D
1 e9 @) R' x8 g, @4 O def __len__(self):3 @0 i x& {6 @% q4 q: [
return len(self.features)
# i3 {2 S( c. z; c0 v3.2 读取数据集. \9 Z. E4 ?/ ?
# \; K+ l8 H& j5 C+ @: V$ U
通过自定义的VOCSegDataset类来分别创建训练集和测试集的实例。因为待会用的是全卷积网络,所以随机裁剪的输出图像的形状可以自己指定,这里指定为320×480 320\times 480320×480。
, r) t; @; a) x& k$ ^* }
, [0 s' v+ b) y; H2 S! obatch_size = 32 # 实际上我的小笔记本不允许我这么做!哭了(大家根据自己电脑内存改吧)
) {5 [ ]* R6 ecrop_size = (320, 480) # 指定随机裁剪的输出图像的形状为(320,480)" U% {5 |" F8 ], r+ E [
max_num = 20000 # 最多从本地读多少张图片,我指定的这个尺寸过滤完不合适的图像之后也就只有1175张~0 W' P6 E6 t6 O/ W' j" |5 E4 y& y
o0 W" [) Q' s x; I, s
# 创建训练集和测试集的实例
+ p# k' y: K" m" F% i/ Uvoc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label, max_num)
7 w7 G% } J8 [1 _; Ivoc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label, max_num)
3 r4 _4 y' a( o0 g* q4 [+ E- Q
3 B; K @3 }( d6 x" [# 设批量大小为32,分别定义【训练集】和【测试集】的数据迭代器
) K Q4 g! c) f/ }4 Y, [* `num_workers = 0 if sys.platform.startswith('win32') else 4
0 h9 g4 k4 N6 G( Y. R) k' `train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
6 H* c' `& K' V- u/ q% Y( p drop_last=True, num_workers=num_workers)
. F/ d% \/ E8 V. ^) qtest_iter = torch.utils.data.DataLoader(voc_test, batch_size, drop_last=True,
0 I; Q" @4 S+ w8 w. j num_workers=num_workers)
: D2 Z* n6 F# k5 a9 O0 i4 h0 x( ~; O# m2 B$ K
# 方便封装,把训练集和验证集保存在dict里
R% n; S# H, a+ h' U& }/ sdataloaders = {'train':train_iter, 'val':test_iter}
+ N* p4 v; D/ G" q& i" ]dataset_sizes = {'train':len(voc_train), 'val':len(voc_test)}" s- d7 H" z$ |( l) z2 i
1 Y2 s3 Q; W9 l9 w' j Z+ I3 q1 l3 N4 构造模型4.1 预训练模型下⾯我们使⽤⼀个基于 ImageNet 数据集预训练的 ResNet-18 模型来抽取图像特征。 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
$ M( C% ]2 b/ N) a8 R
F, @% |7 {7 V$ z" g- xnum_classes = 21 # 21分类,1个背景,20个物体
- t. y# S6 ^" N4 [model_ft = resnet18(pretrained=True) # 设置True,表明要加载使用训练好的参数! v/ \8 s) ]/ f. \) e- Z
8 [4 @* G1 b0 ~. w
# 特征提取器, G6 n$ E8 l& ]7 r8 }. i
for param in model_ft.parameters():! S& ^! {: u m( c% a, ?
param.requires_grad = False O% V, W7 t5 I
4.2 修改成FCN( Q0 f! Y! c! O2 L4 }, E: E& b
! [$ H- b0 B. y
全卷积⽹络(顾名思义全部都是卷积层)先使⽤卷积神经⽹络抽取图像特征,然后通过 1×1 1\times 11×1 卷积层将通道数变换为类别个数,最后通过转置卷积层将特征图的⾼和宽变换为输⼊图像的尺⼨。模型输出与输⼊图像的⾼和宽相同,并在空间位置上⼀⼀对应:3 ]( H. R& N3 a, j$ d
最终输出的通道包含了该空间位置像素的类别预测。* i2 F, ? {; b; M8 ]
) B" U) l' Q8 K( R$ g6 N j对于转置卷积层,如果步幅为 S SS、填充为 S/2 S/2S/2 (假设为整数)、卷积核的⾼和宽为 2S 2S2S,转置卷积核将输⼊的⾼和宽分别放⼤ S SS 倍。
$ L5 x |* ?/ e% ]' X& E$ Q( s4 {" _" Q
可以先打印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 倍即可。
( ~/ ?4 b, t- C: e! f- x% W( W$ G* |7 s& @; H$ u8 m
model_ft = nn.Sequential(*list(model_ft.children())[:-2], # 去掉最后两层# Y8 M" l( C+ l J, }! O
nn.Conv2d(512,num_classes,kernel_size=1), # 用大小为1的卷积层改变输出通道为num_class+ v6 k4 P7 ], t) Z! e# h" W
nn.ConvTranspose2d(num_classes,num_classes, kernel_size=64, padding=16, stride=32)).to(device) # 转置卷积层使图像变为输入图像的大小! p; _* t: E: A; b& f1 e2 X
, W7 q' P: s" p, A$ ^
# 对model_ft做一个测试
# e3 ]- [1 K6 V* E0 q5 u9 Sx = torch.rand((2,3,320,480), device=device) # 构造随机的输入数据) g' F9 b4 c( R, Y4 P% ~& h
print(net(x).shape) # 输出依然是 torch.Size([2, 21, 320, 480])
( [9 ~+ {/ X4 ]+ [& @4 D, w. U- e+ |% r- x1 b0 Z( Y
# 打印第一个小批量的类型和形状。不同于图像分类和目标识别,这里的标签是一个三维数组9 s5 P, W1 _, B/ `. k% C
# for X, Y in train_iter:& L6 d7 K! s5 k# {4 o
# print(X.dtype, X.shape), Q% [+ m. F3 ] X! D
# print(Y.dtype, Y.shape)
+ W& g9 }7 x$ |# break
$ A; o2 e' S9 T: c. X O6 O
: G# T0 @' I+ [8 E* K, h' s D( Z: S. s9 s+ U% i$ Y: V& A) t
4.3 初始化转置卷积层: ?' c, j. x, v2 l" U8 w+ j( h
/ a0 H# b6 u, a. s/ Q在图像处理中,我们有时需要将图像放⼤,即上采样(upsample)。上采样的⽅法有很多,常⽤的有双线性插值。简单来说,为了得到输出图像9 d4 C1 q3 S0 Y* U* Y8 {! d
在坐标 (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函数构造的卷积核的转置卷积层来实现。1 |0 _$ q# |8 ?' C$ F: B$ d
( q+ w- C3 z5 U8 N! Q0 L
/ m6 }+ o8 m8 B6 [# Z+ ?# 双线性插值的上采样,用来初始化转置卷积层的卷积核9 G& Q" i5 @) ] m
def bilinear_kernel(in_channels, out_channels, kernel_size):
+ R0 A7 ?; u9 u- Q! @ factor = (kernel_size+1)//28 ~% y6 F8 h2 v3 D4 l5 ]7 D
if kernel_size%2 == 1:4 V: s7 B+ A6 r
center = factor-1; n5 e' E- R% {% }2 l
else:7 P0 v( g7 S: J$ A
center = factor-0.5
" D2 a/ g/ V* l! J9 H& O og = np.ogrid[:kernel_size, :kernel_size]2 n: H& U* t. j: V% q! h7 J
filt = (1-abs(og[0]-center)/factor) * (1-abs(og[1]-center)/factor)5 g: Z: P; C) f7 `9 Y! G8 `
weight = np.zeros((in_channels,out_channels, kernel_size,kernel_size), dtype='float32'), p5 t; d. U$ ]/ {+ L
weight[range(in_channels), range(out_channels), :, :] = filt
4 { y1 H. S) T weight = torch.Tensor(weight)
, d. ]5 A k0 L u- Z0 f! b weight.requires_grad = True }' P0 l) Y* _) D
return weight
% B% M; H" }/ K- g$ N M
) h2 f N6 b3 c( O
6 G* z }6 G( \; I. P在全卷积⽹络中,将转置卷积层初始化为双线性插值的上采样。对于1×1 1\times 11×1卷积层,采⽤Xavier XavierXavier随机初始化。
) Z3 b$ q g) X# q+ @: @# q! o L( h7 @8 a! E% D7 \! ~3 q3 y
nn.init.xavier_normal_(model_ft[-2].weight.data, gain=1)& _4 [8 m5 w# {( f; @
model_ft[-1].weight.data = bilinear_kernel(num_classes, num_classes, 64).to(device)2 b$ ^! [2 I1 h# }
* a: @/ @7 `+ | K
8 G+ R8 ]/ R, _& h4 O. @; q- X9 W
8 J2 v0 @' m3 C5 训练模型现在可以开始训练模型了。这⾥的损失函数和准确率计算与图像分类中的并没有本质上的不同。有一个 blog 我认为说的很详细,图也画得很好:https://blog.csdn.net/Fcc_bd_stars/article/details/105158215 # Z( `+ B# ]8 k5 r
def train_model(model:nn.Module, criterion, optimizer, scheduler, num_epochs=20):
3 ]. s2 |7 w* @+ Z4 O/ t3 I since = time.time()
! C- Y+ O3 C0 L/ [" r9 x best_model_wts = copy.deepcopy(model.state_dict())' r; O7 {$ N% c3 f4 y8 P
best_acc = 0.0
. k, V$ w5 | i4 i& e0 Z: B' q # 每个epoch都有一个训练和验证阶段- H" V1 o2 ?' |% |% _1 [+ {7 B- B
for epoch in range(num_epochs):
8 j4 `6 u7 M9 z) e3 O print('Epoch {}/{}'.format(epoch, num_epochs-1))
/ x; H3 I. n" {) k/ ~$ N print('-'*10)
9 @( i7 u2 v9 }6 f( u% \ for phase in ['train', 'val']:3 y U; P6 l! Y4 e t
if phase == 'train':8 o+ i# O) P4 X; V. m3 X, V
scheduler.step()
, Q$ V4 U* t5 m9 G: I9 Q y model.train(); n9 Y: t& p, x! o1 }9 a6 n$ ~
else:/ t6 t# E/ W" e1 p' M
model.eval(): Y3 w$ p3 w# V! ^
runing_loss = 0.0/ M) ~+ D F5 G% n/ r: i9 U% O
runing_corrects = 0.0
4 L( L9 c$ b p; I: Y1 H # 迭代一个epoch
( ^" C5 M3 m$ P. j: g! i2 Y for inputs, labels in dataloaders[phase]:: m. u' ^, {" O9 J( s
inputs, labels = inputs.to(device), labels.to(device)
: ]4 c# J; j7 w9 d0 q# L5 j optimizer.zero_grad() # 零参数梯度
( v y6 l4 n' b( u8 y # 前向,只在训练时跟踪参数$ s: w1 x, ?% E9 ]
with torch.set_grad_enabled(phase=='train'):
% U$ Q+ T; Y7 M. h0 B logits = model(inputs) # [5, 21, 320, 480]
# Q& d8 }) H- L& L: N1 u# G loss = criteon(logits, labels.long())
3 c+ |% t- n v& t7 S7 J* ^) c! T # 后向,只在训练阶段进行优化& t( l& L; G' N1 ^6 t
if phase=='train':
7 K" e- p- A: h) I% N loss.backward()9 Z& X, [$ l0 a
optimizer.step()8 y" K- H9 }8 x1 D; |" S M
# 统计loss和correct
& a, w- c9 R" D9 k) K runing_loss += loss.item()*inputs.size(0)7 k; A! ~4 b" @
runing_corrects += torch.sum((torch.argmax(logits.data,1))==labels.data)/(480*320)
, P% I2 c8 a; A4 x3 X- H' k0 [- e9 E R
epoch_loss = runing_loss / dataset_sizes[phase]
5 n- d2 ~ \; M; J5 ~ epoch_acc = runing_corrects.double() / dataset_sizes[phase]
+ N w D t$ r5 u, X/ ? k, c print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
/ i7 K! i5 b$ N. M2 h # 深度复制model参数
3 s9 B' V- R1 N0 u if phase=='val' and epoch_acc>best_acc:
5 r; o0 v) f* n+ J4 n1 ] best_acc = epoch_acc
/ ~5 F6 J2 V5 ]7 [ best_model_wts = copy.deepcopy(model.state_dict())
! |; }- I9 f2 E/ u( s print()2 {. E5 N& @; N! u' S
time_elapsed = time.time() - since;2 S+ D4 s! J9 {3 T8 r7 Z
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
3 p3 ]' M: w& B2 m # 加载最佳模型权重
: m9 D" S: B. w1 c model.load_state_dict(best_model_wts)
* \$ D5 d3 N3 p% E* m return model
8 ?# A5 ]7 l9 u& h. X- a5 |8 _2 ?7 f
6 t! f1 y( j$ f+ n% J, `下面定义train_model要用到的参数,开始训练! U6 x* O6 G2 _7 L
* u0 i3 b4 R, K/ N
epochs = 5 # 训练5个epoch7 C+ q( V1 {- b. Q4 B( y
criteon = nn.CrossEntropyLoss()
1 }9 j# H8 Z3 C% u8 M7 J' G. Joptimizer = optim.SGD(model_ft.parameters(), lr=0.001, weight_decay=1e-4, momentum=0.9)
1 b+ F2 j1 P, h. n7 x* y# 每3个epochs衰减LR通过设置gamma=0.1
3 |2 S; [( r3 P5 W+ Z" b4 ?: `5 Rexp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
& t: P. R, i8 j: w6 p9 T: z
; p5 R/ z) U4 ~' u* `7 l& k' h# 开始训练
4 b+ k* P9 B+ b: gmodel_ft = train_model(model_ft, criteon, optimizer, exp_lr_scheduler, num_epochs=epochs)- q8 u; x% r6 e8 W
5 E( V: T$ w6 t% o6 测试模型为了可视化每个像素的预测类别,我们将预测类别映射回它们在数据集中的标注颜⾊。 def label2image(pred):4 q" _# c! [* h9 K ?
# pred: [320,480]
. g' v9 I. [3 \& @, t/ Y+ D colormap = torch.tensor(VOC_COLORMAP,device=device,dtype=int)
$ o3 X7 q* L! p) i, n: N x = pred.long()
1 H. { E- r$ Q/ n- G return (colormap[x,:]).data.cpu().numpy()5 D _3 O e8 `6 A
" e$ N" |7 t' n* t: ~9 F( [0 h
下面这里提供了两种测试形式 6.1 通用型其实如果要用于测试其它数据集,也是要改动一下的 : ) 😃 mean=torch.tensor([0.485, 0.456, 0.406]).reshape(3,1,1).to(device)
3 T. D' s5 B, j! {1 E& [4 W0 Qstd=torch.tensor([0.229, 0.224, 0.225]).reshape(3,1,1).to(device)* U: n, g' W6 q4 b- ?2 Q* v
def visualize_model(model:nn.Module, num_images=4):$ f& F- d+ w4 `; E7 z! w
was_training = model.training
9 p1 q* F, |' Y- T3 Q) _/ y model.eval()) z% N8 H' u! R" y7 _& Q, }- F
images_so_far = 0
" s( N0 J$ J( `2 [9 h, ]7 g n, imgs = num_images, []
5 k a9 d5 o2 a5 w' e6 K9 \ with torch.no_grad():% q& T+ s0 y* C3 }0 v
for i, (inputs, labels) in enumerate(dataloaders['val']):
. \" c$ l1 c9 c2 f+ t$ m inputs, labels = inputs.to(device), labels.to(device) # [b,3,320,480]
+ P: h3 W3 g2 v4 J outputs = model(inputs)
: q6 Z2 c0 }8 h" Z4 Q pred = torch.argmax(outputs, dim=1) # [b,320,480]
! k$ Q; K8 G" J; Y, ^% J4 n inputs_nd = (inputs*std+mean).permute(0,2,3,1)*255 # 记得要变回去哦
8 ` a% Q0 v- m6 f
( [: [6 ^% J' E& [ for j in range(num_images):
% X# v& x( M9 H: x images_so_far += 1( F" q0 S. {* I: W1 F+ J" D
pred1 = label2image(pred[j]) # numpy.ndarray (320, 480, 3)
' Q7 j1 _5 |; H, w* _- R$ ~ imgs += [inputs_nd[j].data.int().cpu().numpy(), pred1, label2image(labels[j])]( o4 h0 p2 |+ @& ^
if images_so_far == num_images:+ m( U( k, F' I l! U
model.train(mode=was_training)
j, G. _$ o+ k% r # 我已经固定了每次只显示4张图了,大家可以自己修改4 A. V( s1 ]$ x* [
show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
) I9 _9 n! ^3 ?6 a. g return model.train(mode=was_training). `; O% y! P! W3 Q0 B# i
" [) d( e* Y$ R/ w' v7 O& [# 开始验证
2 ^! R8 h+ K* R7 B( Q, P4 l7 @2 fvisualize_model(model_ft)9 Y! W7 d+ D3 C
6.2 不通用在预测时,我们需要将输⼊图像在各个通道做标准化,并转成卷积神经⽹络所需要的四维输⼊格式。 # 预测前将图像标准化,并转换成(b,c,h,w)的tensor
6 A" H. D! W, f, Y2 Ndef predict(img, model):6 H& _( ]0 ], V) `
tsf = transforms.Compose([) e* u4 O$ i& X8 x6 V
transforms.ToTensor(), # 好像会自动转换channel
% _$ E# }& x5 m% [ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
: b' J( b4 f, X& L x = tsf(img).unsqueeze(0).to(device) # (3,320,480) -> (1,3,320,480)
" `' l) Y( ~0 a pred = torch.argmax(model(x), dim=1) # 每个通道选择概率最大的那个像素点 -> (1,320,480)/ b& t& _" i8 X6 `
return pred.reshape(pred.shape[1],pred.shape[2]) # reshape成(320,480)
0 y A ^3 ~$ B4 \ r$ t
5 {( o( Z- S3 b( g2 e# cdef evaluate(model:nn.Module):
1 w: [7 j; V! T model.eval()
$ u+ R3 q& }# g test_images, test_labels = read_voc_images(voc_dir, is_train=False, max_num=10) " _. O7 m/ s0 [
n, imgs = 4, []
' v2 f- m6 G D& b9 Y2 Q: u for i in range(n):
9 p! K% Z) W* s$ `" i5 Z xi, yi = voc_rand_crop(test_images, test_labels, 320, 480) # Image5 u% _$ x+ h$ N. J; m2 z n
pred = label2image(predict(xi, model))
/ c+ v& O. p) S( R) y! F6 M; M imgs += [xi, pred, yi]& w: V+ _' i, _" h
show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n)
; Y4 f0 _1 @$ J$ Z, @$ }( H/ Z' j+ F3 J
# 开始测试
. T: q2 e/ |/ Bevaluate(model_ft)/ L$ s9 c% t0 j% M% N$ a/ ]
: ]3 B3 k1 p6 b, D0 a5 K/ |7 结语我只训练了3个epoch,下面是训练输出
1 b% ?9 J) s7 wEpoch 0/2
- g& o) R; H/ c, t* T----------
# ~: Z Z7 `' g* Ftrain Loss: 1.7844 Acc: 0.5835
/ _, z- }+ W" X6 R aval Loss: 1.1669 Acc: 0.64563 U( z5 e& v! ^7 I6 }) |$ b$ e, J
% R& g' N9 ]% l# rEpoch 1/2( \2 [5 j. S. n3 S1 D6 D
----------
& N, m9 `* v. Vtrain Loss: 1.1288 Acc: 0.6535* r' m$ B& w% w
val Loss: 0.9012 Acc: 0.6929
3 |+ W C( n4 j* d/ T+ {) X* u' ]; |: u; L5 L* b
Epoch 2/28 o" U! B3 D' S
----------
+ F% I9 o( Q, L. L; X0 itrain Loss: 0.9578 Acc: 0.67063 _( t) n6 T3 O" {5 P9 J! o$ j
val Loss: 0.8088 Acc: 0.6948
4 C5 D& }* |/ h' g# ?5 L; A
- J" Z8 }1 l, g. X5 M6 G) M0 PTraining complete in 6m 37s
% t' C3 }3 \3 r7 v; S/ b" w# C' y8 H
; m$ W% [# S( b1 o9 M
- |0 t( O( D/ {; A% s' D W! `
; F7 T: c! r* d8 M9 b8 T3 S3 e0 O当 epochs = 5 时,训练集的精度在 89 8989% 左右,测试集的精度可以达到 86 8686 %。5 T5 W2 e# ]. i6 z6 q; w
( S f/ z$ a$ s0 L/ p' d* m
对于这个模型用 ResNet-50 作特征提取器会有更好的效果,不过训练的时间也会更长。还有超参数lr, weight_decay, momentum, step_size, gamma 以及1×1 1×11×1卷积层和转置卷积层的初始化方式也可以继续调。
) G/ {6 ~1 y% M4 \/ x& C+ A3 ^. U1 K/ Y) Z
/ e: _( T& B) u9 P7 j5 B2 e+ n
语义分割还有很多可用的模型,本文用的是 FCN,在其它一些模型上会有更好的表现:
- B7 f- ]' ~5 b; e7 L: f5 W9 O) }3 z4 p1 T
Deeplab V3+ 具有可分离卷积的编码器/解码器,用于语义图像分割[论文]! f/ C8 p$ {# K% }. I7 ~' B+ l
GCN 通过全局卷积网络改进语义分割[论文]# X3 s- F* N4 S
UperNet 统一感知解析& ^- ~/ U6 _* F# b& ~0 @5 [; _4 @
ENet 用于实时语义分割的深度神经网络体系结构[论文]
& j, T5 h( P9 sU-Net 用于生物医学图像分割的卷积网络5 n( R @: x D6 s
SegNet 用于图像分段的深度卷积编码器-解码器架构。
% ?# ^7 A( z1 Z9 b+ Y还有(DUC,HDC)、PSPNet等。
2 F* E" y; w! G9 I+ X, k6 j% p3 N$ J* O& N" r/ t4 _
常用的语义分割数据集也有很多:Pascal VOC、CityScapes、ADE20K、COCO Stuff等。& k) m: _# U+ T' ^7 B/ S
6 O6 C. R5 X- b8 z/ T0 b; n2 H/ n对于损失函数,除了交叉熵误差,也可以用这些:
+ M2 Q: c! J, c' H! H
% L- e! q6 h, A% ^Dice-Loss 可以测试两个样本之间的重叠度量,可以更好地反映训练目标,但该损失函数具有很强的非凸性,很难优化。
( |8 Y- S9 C( `& [% mCE Dice loss Dice 损失与 CE 的总和,CE 提供了平滑的优化,而 Dice 损失则很好地表明了分割结果的质量。1 \8 H' \5 M6 S# p% i8 D
Focal Loss CE 的另一种版本,用于避免类别不平衡而降低了置信度的情况。9 v5 N8 X: c' h* z1 P( ^
Lovasz Softmax 查看论文:Lovasz - softmax损失。+ k) K! y$ k# n8 b) u% @7 y
; \8 Q! |3 k+ i+ J
0 X- o1 x& X% U, {2 T* ^: \
: @# z: H- I y) o" g
- U' o9 t1 e. k! y; i
% b% ?2 ]5 r+ Y————————————————- f/ G5 f) x( {2 v
版权声明:本文为CSDN博主「小红不吃糖」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。, o$ ~' }% H9 k+ j3 C% `' I
原文链接:https://blog.csdn.net/qq_43280818/article/details/105916507
) G) q" S$ i8 R) r/ i. p: A! ~, B% t/ o5 S! {+ A$ ~$ K- {
3 O6 G9 b1 n+ v+ ]) K- Y4 H3 ?' y
|