Caffe 中的上采样操作

官方 Caffe 中没有实现上采样层, 但提供了一种迂回的实现方法, 即采用特殊参数配置的 Deconvolution 层 (Deconvolution, 直译为反卷积, 但更一般的称呼是转置卷积). 下面是测试代码:

import os
import caffe
import numpy as np
from caffe import layers as L
from caffe.proto import caffe_pb2
import skimage


def create_upsample_net(shape, factor):
    kernel_size = 2 * factor - factor % 2
    pad = factor // 2
    stride = factor
    num_output = shape[1]
    group = shape[1]

    blob_shape = caffe_pb2.BlobShape()
    blob_shape.dim.extend(shape)
    n = caffe.NetSpec()
    n.data = L.Input(shape=[blob_shape])
    convolution_param = dict(num_output=num_output, 
                             kernel_size=kernel_size, 
                             stride=stride, 
                             group=num_output,
                             pad=pad,
                             weight_filler=dict(type="bilinear"),
                             bias_term=False)
    n.conv = L.Deconvolution(n.data, 
                             convolution_param=convolution_param,
                             param=dict(lr_mult=0, decay_mult=0))

    prototxt = 'temp.prototxt'
    with open(prototxt, 'w') as f:
        f.write(str(n.to_proto()))
    net = caffe.Net(prototxt, caffe.TEST)
    return net


def test_caffe_upsample(data, factor):
    net = create_upsample_net(data.shape, factor)
    net.blobs['data'].data[...] = data
    net.forward()
    caffe_output = net.blobs['conv'].data

    x = np.squeeze(data, axis=0)
    x = np.transpose(x, (1,2,0))
    skimage_output = skimage.transform.rescale(x, factor, mode='constant', cval=0)
    skimage_output = np.transpose(skimage_output, (2,0,1))
    skimage_output = np.expand_dims(skimage_output, axis=0)

    print('MAE: {:.7f}'.format(np.mean(np.abs(caffe_output - skimage_output))))
    print('MAX: {:.7f}'.format(np.max(np.abs(caffe_output - skimage_output))))
    print(caffe_output.flatten()[5:10])
    print(skimage_output.flatten()[5:10])


if __name__ == '__main__':
    np.random.seed(0)
    x = np.random.randn(1, 3, 112, 112)
    test_caffe_upsample(x, factor=2)

经测试, Deconvolution 实现的 (双线性) 上采样操作的结果与 skimage.transform.rescale 的结果几乎完全一致.

用转置卷积实现上采样的参数设置说明

1): num_output 和 group 等于输入通道数 (shape[1]), 因为上采样操作不会改变 feature map 的通道数, 而且上采样输出 feature map 的某个通道只与输入 feature map 的对应通道有关.

2): kernel_size, pad 和 stride 的设置, 可以参考下面的推导:

在 $CAFFE_ROOT/src/caffe/layers/deconv_layer.cpp 文件中

const int kernel_extent = dilation_data[i] * (kernel_shape_data[i] - 1) + 1;
const int output_dim = stride_data[i] * (input_dim - 1) + 
    kernel_extent - 2 * pad_data[i];

不妨将上面代码转化为数学语言: 设输入尺寸为 $N$, 卷积核尺寸为 $K$, 步长为 $S$, 填充尺寸为 $P$, 膨胀系数为 $D$, 则输出尺寸 $M$ 为:
$$ E = D\times (K - 1) + 1 $$
$$ M = S \times (N - 1) + E - 2P $$

特别地, 当 $D=1$ 时, 上式简化为为:
$$ M = S \times (N - 1) + K - 2P $$

又当 $S = f$, $K = 2\times f - f \% 2$, $P=\lfloor f / 2 \rfloor$, 其中 $f$ 是正整数 (可与上面的代码对应).

$$\begin{aligned} M &= f\times (N - 1) + 2\times f - f \% 2 - 2 \lfloor f / 2 \rfloor \\ &= f\times (N - 1) + 2\times f - f \\ &= f\times N \\ \end{aligned}$$

可见输出尺寸 $M$ 是输入尺寸 $N$ 的 $f$ 倍.

3): bilinear 是 Caffe 的一种参数初始化方式, 顾名思义, 就是将卷积核初始化为双线性插值的系数; 因为双线性插值是线性运算, 所以 bias_tem 需设置为 False.

4): lr_mult 和 decay_mult 要设置为 0, 因为要保证该层一直执行的是双线性插值, 在训练中该层的卷积核权重不能变.

与 PyTorch 中的双线性插值上采样的比较

笔者还将 Caffe 中的上采样操作 (为了简便, 实际上用的是 skimage.transform.rescale) 与 PyTorch 中的 torch.nn.Upsample (mode 为 bilinear, align_corners 为 True 和 False) 进行了比较, 很遗憾, 结果显示不一致, 测试代码如下.

import cv2
import numpy as np
import skimage
import skimage.transform

import torch


def test_pytorch_upsample(data, factor):
    x = np.squeeze(data, axis=0)
    x = np.transpose(x, (1,2,0))
    skimage_output = skimage.transform.rescale(x, factor, mode='constant', 
                                               cval=0, multichannel=True)
    skimage_output = np.transpose(skimage_output, (2,0,1))
    skimage_output = np.expand_dims(skimage_output, axis=0)

    upsample = torch.nn.Upsample(scale_factor=factor, mode='bilinear', 
                                 align_corners=False)
    torch_output = upsample(torch.from_numpy(data)).numpy()

    upsample_2 = torch.nn.Upsample(scale_factor=factor, mode='bilinear', 
                                   align_corners=True)
    torch_output_2 = upsample_2(torch.from_numpy(data)).numpy()

    print('MAE: {:.7f}'.format(np.mean(np.abs(torch_output - skimage_output))))
    print('MAX: {:.7f}'.format(np.max(np.abs(torch_output - skimage_output))))
    print('MAE: {:.7f}'.format(np.mean(np.abs(torch_output_2 - skimage_output))))
    print('MAX: {:.7f}'.format(np.max(np.abs(torch_output_2 - skimage_output))))
    print(skimage_output.flatten()[5:10])
    print(torch_output.flatten()[5:10])
    print(torch_output_2.flatten()[5:10])


if __name__ == '__main__':
    np.random.seed(0)
    data = np.random.randn(1, 3, 112, 112).astype(np.float32)
    test_pytorch_upsample(data, factor=2)

在 Caffe 和 PyTorch 模型格式转换时可能会踩到这个坑, 所以如果你有 "Caffe 和 PyTorch 之间模型格式转换的需求", 为了保证较高保真的转换, 最好在网络设计阶段就避免使用上采样层, 或者能在 Caffe 或 PyTorch 中实现结果一致的上采样层. 下面提供一种方案: 在 Caffe 中不使用 Deconvolution 实现的 (双线性插值的) 上采样层, 在 PyTorch 中也不使用 torch.nn.Upsample, 而是统一都采用一般的转置卷积. 对于 PyTorch, 则有:

class Upsample(nn.Module):
    def __init__(self, c1, c2=None, k=2):
        super(Upsample, self).__init__()
        c2 = c2 if c2 is not None else c1
        kernel_size = 2 * k - k % 2
        pad = k // 2
        stride = k
        self.conv_trans = nn.ConvTranspose2d(c1, c2, kernel_size, stride=stride, 
                                             padding=pad, output_padding=0,  
                                             groups=c1, bias=False)

    def forward(self, x):
        return self.conv_trans(x)

对于 Caffe 则有:

kernel_size = 2 * factor - factor % 2
pad = factor // 2
stride = factor
num_output = shape[1]
group = shape[1]

convolution_param = dict(num_output=num_output, 
                         kernel_size=kernel_size, 
                         stride=stride, 
                         group=group,
                         pad=pad,
                         bias_term=False)

注意没有了 weight_filler, 同时如果是训练的话, lr_mult 和 decay_mult 还需设置成合适的值.

更新记录

20210708, 发布

版权声明

保持署名-非商业用途-非衍生, 知识共享3.0协议.