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