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