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