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