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