全卷积网络(FCN)

优质
小牛编辑
113浏览
2023-12-01

上一节介绍了,我们可以基于语义分割对图像中的每个像素进行类别预测。全卷积网络(fully convolutional network,FCN)采用卷积神经网络实现了从图像像素到像素类别的变换 [1]。与之前介绍的卷积神经网络有所不同,全卷积网络通过转置卷积(transposed convolution)层将中间层特征图的高和宽变换回输入图像的尺寸,从而令预测结果与输入图像在空间维(高和宽)上一一对应:给定空间维上的位置,通道维的输出即该位置对应像素的类别预测。

我们先导入实验所需的包或模块,然后解释什么是转置卷积层。

%matplotlib inline
import d2lzh as d2l
from mxnet import gluon, image, init, nd
from mxnet.gluon import data as gdata, loss as gloss, model_zoo, nn
import numpy as np
import sys

转置卷积层

顾名思义,转置卷积层得名于矩阵的转置操作。事实上,卷积运算还可以通过矩阵乘法来实现。在下面的例子中,我们定义高和宽分别为4的输入X,以及高和宽分别为3的卷积核K。打印二维卷积运算的输出以及卷积核。可以看到,输出的高和宽分别为2。

X = nd.arange(1, 17).reshape((1, 1, 4, 4))
K = nd.arange(1, 10).reshape((1, 1, 3, 3))
conv = nn.Conv2D(channels=1, kernel_size=3)
conv.initialize(init.Constant(K))
conv(X), K

下面我们将卷积核K改写成含有大量零元素的稀疏矩阵W,即权重矩阵。权重矩阵的形状为(4, 16),其中的非零元素来自卷积核K中的元素。将输入X逐行连结,得到长度为16的向量。然后将W与向量化的X做矩阵乘法,得到长度为4的向量。对其变形后,我们可以得到和上面卷积运算相同的结果。可见,我们在这个例子中使用矩阵乘法实现了卷积运算。

W, k = nd.zeros((4, 16)), nd.zeros(11)
k[:3], k[4:7], k[8:] = K[0, 0, 0, :], K[0, 0, 1, :], K[0, 0, 2, :]
W[0, 0:11], W[1, 1:12], W[2, 4:15], W[3, 5:16] = k, k, k, k
nd.dot(W, X.reshape(16)).reshape((1, 1, 2, 2)), W

现在我们从矩阵乘法的角度来描述卷积运算。设输入向量为$\boldsymbol{x}$,权重矩阵为$\boldsymbol{W}$,卷积的前向计算函数的实现可以看作将函数输入乘以权重矩阵,并输出向量$\boldsymbol{y} = \boldsymbol{W}\boldsymbol{x}$。我们知道,反向传播需要依据链式法则。由于$\nabla_{\boldsymbol{x}} \boldsymbol{y} = \boldsymbol{W}^\top$,卷积的反向传播函数的实现可以看作将函数输入乘以转置后的权重矩阵$\boldsymbol{W}^\top$。而转置卷积层正好交换了卷积层的前向计算函数与反向传播函数:这两个函数可以看作将函数输入向量分别乘以$\boldsymbol{W}^\top$和$\boldsymbol{W}$。

不难想象,转置卷积层可以用来交换卷积层输入和输出的形状。让我们继续用矩阵乘法描述卷积。设权重矩阵是形状为$4\times16$的矩阵,对于长度为16的输入向量,卷积前向计算输出长度为4的向量。假如输入向量的长度为4,转置权重矩阵的形状为$16\times4$,那么转置卷积层将输出长度为16的向量。在模型设计中,转置卷积层常用于将较小的特征图变换为更大的特征图。在全卷积网络中,当输入是高和宽较小的特征图时,转置卷积层可以用来将高和宽放大到输入图像的尺寸。

我们来看一个例子。构造一个卷积层conv,并设输入X的形状为(1, 3, 64, 64)。卷积输出Y的通道数增加到10,但高和宽分别缩小了一半。

conv = nn.Conv2D(10, kernel_size=4, padding=1, strides=2)
conv.initialize()

X = nd.random.uniform(shape=(1, 3, 64, 64))
Y = conv(X)
Y.shape

下面我们通过创建Conv2DTranspose实例来构造转置卷积层conv_trans。这里我们设conv_trans的卷积核形状、填充以及步幅与conv中的相同,并设输出通道数为3。当输入为卷积层conv的输出Y时,转置卷积层输出与卷积层输入的高和宽相同:转置卷积层将特征图的高和宽分别放大了2倍。

conv_trans = nn.Conv2DTranspose(3, kernel_size=4, padding=1, strides=2)
conv_trans.initialize()
conv_trans(Y).shape

在有些文献中,转置卷积也被称为分数步长卷积(fractionally-strided convolution)[2]。

构造模型

我们在这里给出全卷积网络模型最基本的设计。如图9.11所示,全卷积网络先使用卷积神经网络抽取图像特征,然后通过$1\times 1$卷积层将通道数变换为类别个数,最后通过转置卷积层将特征图的高和宽变换为输入图像的尺寸。模型输出与输入图像的高和宽相同,并在空间位置上一一对应:最终输出的通道包含了该空间位置像素的类别预测。

下面我们使用一个基于ImageNet数据集预训练的ResNet-18模型来抽取图像特征,并将该网络实例记为pretrained_net。可以看到,该模型成员变量features的最后两层分别是全局最大池化层GlobalAvgPool2D和样本变平层Flatten,而output模块包含了输出用的全连接层。全卷积网络不需要使用这些层。

pretrained_net = model_zoo.vision.resnet18_v2(pretrained=True)
pretrained_net.features[-4:], pretrained_net.output

下面我们创建全卷积网络实例net。它复制了pretrained_net实例成员变量features里除去最后两层的所有层以及预训练得到的模型参数。

net = nn.HybridSequential()
for layer in pretrained_net.features[:-2]:
    net.add(layer)

给定高和宽分别为320和480的输入,net的前向计算将输入的高和宽减小至原来的$1/32$,即10和15。

X = nd.random.uniform(shape=(1, 3, 320, 480))
net(X).shape

接下来,我们通过$1\times 1$卷积层将输出通道数变换为Pascal VOC2012数据集的类别个数21。最后,我们需要将特征图的高和宽放大32倍,从而变回输入图像的高和宽。回忆一下“填充和步幅”一节中描述的卷积层输出形状的计算方法。由于$(320-64+16\times2+32)/32=10$且$(480-64+16\times2+32)/32=15$,我们构造一个步幅为32的转置卷积层,并将卷积核的高和宽设为64、填充设为16。不难发现,如果步幅为$s$、填充为$s/2$(假设$s/2$为整数)、卷积核的高和宽为$2s$,转置卷积核将输入的高和宽分别放大$s$倍。

num_classes = 21
net.add(nn.Conv2D(num_classes, kernel_size=1),
        nn.Conv2DTranspose(num_classes, kernel_size=64, padding=16,
                           strides=32))

初始化转置卷积层

我们已经知道,转置卷积层可以放大特征图。在图像处理中,我们有时需要将图像放大,即上采样(upsample)。上采样的方法有很多,常用的有双线性插值。简单来说,为了得到输出图像在坐标$(x,y)$上的像素,先将该坐标映射到输入图像的坐标$(x',y')$,例如,根据输入与输出的尺寸之比来映射。映射后的$x'$和$y'$通常是实数。然后,在输入图像上找到与坐标$(x',y')$最近的4个像素。最后,输出图像在坐标$(x,y)$上的像素依据输入图像上这4个像素及其与$(x',y')$的相对距离来计算。双线性插值的上采样可以通过由以下bilinear_kernel函数构造的卷积核的转置卷积层来实现。限于篇幅,我们只给出bilinear_kernel函数的实现,不再讨论算法的原理。

def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
                      dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
    return nd.array(weight)

我们来实验一下用转置卷积层实现的双线性插值的上采样。构造一个将输入的高和宽放大2倍的转置卷积层,并将其卷积核用bilinear_kernel函数初始化。

conv_trans = nn.Conv2DTranspose(3, kernel_size=4, padding=1, strides=2)
conv_trans.initialize(init.Constant(bilinear_kernel(3, 3, 4)))

读取图像X,将上采样的结果记作Y。为了打印图像,我们需要调整通道维的位置。

img = image.imread('../img/catdog.jpg')
X = img.astype('float32').transpose((2, 0, 1)).expand_dims(axis=0) / 255
Y = conv_trans(X)
out_img = Y[0].transpose((1, 2, 0))

可以看到,转置卷积层将图像的高和宽分别放大2倍。值得一提的是,除了坐标刻度不同,双线性插值放大的图像和“目标检测和边界框”一节中打印出的原图看上去没什么两样。

d2l.set_figsize()
print('input image shape:', img.shape)
d2l.plt.imshow(img.asnumpy());
print('output image shape:', out_img.shape)
d2l.plt.imshow(out_img.asnumpy());

在全卷积网络中,我们将转置卷积层初始化为双线性插值的上采样。对于$1\times 1$卷积层,我们采用Xavier随机初始化。

net[-1].initialize(init.Constant(bilinear_kernel(num_classes, num_classes,
                                                 64)))
net[-2].initialize(init=init.Xavier())

读取数据集

我们用上一节介绍的方法读取数据集。这里指定随机裁剪的输出图像的形状为$320\times 480$:高和宽都可以被32整除。

crop_size, batch_size, colormap2label = (320, 480), 32, nd.zeros(256**3)
for i, cm in enumerate(d2l.VOC_COLORMAP):
    colormap2label[(cm[0] * 256 + cm[1]) * 256 + cm[2]] = i
voc_dir = d2l.download_voc_pascal(data_dir='../data')

num_workers = 0 if sys.platform.startswith('win32') else 4
train_iter = gdata.DataLoader(
    d2l.VOCSegDataset(True, crop_size, voc_dir, colormap2label), batch_size,
    shuffle=True, last_batch='discard', num_workers=num_workers)
test_iter = gdata.DataLoader(
    d2l.VOCSegDataset(False, crop_size, voc_dir, colormap2label), batch_size,
    last_batch='discard', num_workers=num_workers)

训练模型

现在可以开始训练模型了。这里的损失函数和准确率计算与图像分类中的并没有本质上的不同。因为我们使用转置卷积层的通道来预测像素的类别,所以在SoftmaxCrossEntropyLoss里指定了axis=1(通道维)选项。此外,模型基于每个像素的预测类别是否正确来计算准确率。

ctx = d2l.try_all_gpus()
loss = gloss.SoftmaxCrossEntropyLoss(axis=1)
net.collect_params().reset_ctx(ctx)
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.1,
                                                      'wd': 1e-3})
d2l.train(train_iter, test_iter, net, loss, trainer, ctx, num_epochs=5)

预测像素类别

在预测时,我们需要将输入图像在各个通道做标准化,并转成卷积神经网络所需要的四维输入格式。

def predict(img):
    X = test_iter._dataset.normalize_image(img)
    X = X.transpose((2, 0, 1)).expand_dims(axis=0)
    pred = nd.argmax(net(X.as_in_context(ctx[0])), axis=1)
    return pred.reshape((pred.shape[1], pred.shape[2]))

为了可视化每个像素的预测类别,我们将预测类别映射回它们在数据集中的标注颜色。

def label2image(pred):
    colormap = nd.array(d2l.VOC_COLORMAP, ctx=ctx[0], dtype='uint8')
    X = pred.astype('int32')
    return colormap[X, :]

测试数据集中的图像大小和形状各异。由于模型使用了步幅为32的转置卷积层,当输入图像的高或宽无法被32整除时,转置卷积层输出的高或宽会与输入图像的尺寸有偏差。为了解决这个问题,我们可以在图像中截取多块高和宽为32的整数倍的矩形区域,并分别对这些区域中的像素做前向计算。这些区域的并集需要完整覆盖输入图像。当一个像素被多个区域所覆盖时,它在不同区域前向计算中转置卷积层输出的平均值可以作为softmax运算的输入,从而预测类别。

简单起见,我们只读取几张较大的测试图像,并从图像的左上角开始截取形状为$320\times480$的区域:只有该区域用来预测。对于输入图像,我们先打印截取的区域,再打印预测结果,最后打印标注的类别。

test_images, test_labels = d2l.read_voc_images(is_train=False)
n, imgs = 4, []
for i in range(n):
    crop_rect = (0, 0, 480, 320)
    X = image.fixed_crop(test_images[i], *crop_rect)
    pred = label2image(predict(X))
    imgs += [X, pred, image.fixed_crop(test_labels[i], *crop_rect)]
d2l.show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n);

小结

  • 可以通过矩阵乘法来实现卷积运算。
  • 全卷积网络先使用卷积神经网络抽取图像特征,然后通过$1\times 1$卷积层将通道数变换为类别个数,最后通过转置卷积层将特征图的高和宽变换为输入图像的尺寸,从而输出每个像素的类别。
  • 在全卷积网络中,可以将转置卷积层初始化为双线性插值的上采样。

练习

  • 用矩阵乘法来实现卷积运算是否高效?为什么?
  • 如果将转置卷积层改用Xavier随机初始化,结果有什么变化?
  • 调节超参数,能进一步提升模型的精度吗?
  • 预测测试图像中所有像素的类别。
  • 全卷积网络的论文中还使用了卷积神经网络的某些中间层的输出 [1]。试着实现这个想法。

参考文献

[1] Long, J., Shelhamer, E., & Darrell, T. (2015). Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3431-3440).

[2] Dumoulin, V., & Visin, F. (2016). A guide to convolution arithmetic for deep learning. arXiv preprint arXiv:1603.07285.