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