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