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