以OneFlow为例梳理深度学习框架的那些插值方法

Posted 2021-09-22 just_sort

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了以OneFlow为例梳理深度学习框架的那些插值方法相关的知识，希望对你有一定的参考价值。

0x0. 前言

这篇文章基于自己为OneFlow框架开发interpolate这个Op总结而来，OneFlow的interpolate Op 和 Pytorch的功能一致，都是用来实现插值上采样或者下采样的。在实现这个Op的时候还给Pytorch修复了一个bug并合并到了主仓库，见：https://github.com/pytorch/pytorch/commit/6ab3a210983b7eee417e7cd92a8ad2677065e470。因此OneFlow框架中的interpolate算子和Pytorch中的interpolate算子的功能是完全等价的。这篇文章就以OneFlow中这个算子的实现为例来盘点一下深度学习框架中的那些插值算法。

0x1. doc && interface接口

要了解interpolate算子中的插值算法，首先需要从文档和Python前端接口看起。看一下接口文档，https://oneflow.readthedocs.io/en/master/functional.html?highlight=interpolate 。

这里可以看到OneFlow的interpolate算子用来实现插值上采样或者下采样的功能，支持3-D，4-D，5-D的输入Tensor，然后提供了多种插值的方式应用于不同Shape的输入Tensor。下面再看一下参数列表：

input：输入Tensor。
size：插值后输出Tensor的空间维度的大小，这个spatial size就是去掉Batch，Channel，Depth维度后剩下的值。比如NCHW的spatial size是HW。
scale_factor(float 或者 Tuple[float])：spatial size的乘数，如果是tuple则必须匹配输入数据的大小。
mode(str)：上采样的模式，包含’nearest’ | ‘linear’ | ‘bilinear’ | ‘bicubic’ | ‘trilinear’ | ‘area’。默认是 ‘nearest’。
align_corners(bool)：在几何上，我们将输入和输出的像素视为正方形而不是点。如果设置为True，则输入和输出张量按其角像素的中心点对齐，保留角像素处的值。如果设置为False，则输入和输出张量按其角像素的角点对齐，插值使用边缘值填充来处理边界外值，当scale_factor保持不变时，此操作与输入大小无关。这仅在mode为 ‘linear’ | ‘bilinear’ | ‘bicubic’ | 'trilinear’时有效。默认值是False。（没看懂没关系，下面有一节专门讲解）
recompute_scale_factor(bool)：重新计算用于插值计算的 scale_factor。当 scale_factor 作为参数传递时，它用于计算 output_size。如果 recompute_scale_factor 为 False 或未指定，则传入的 scale_factor 将用于插值计算。否则，将根据用于插值计算的输出和输入大小计算新的 scale_factor（即，等价于显示传入output_size）。请注意，当 scale_factor 是浮点数时，由于舍入和精度问题，重新计算的 scale_factor 可能与传入的不同。

除了功能描述和参数描述之外还有几个注意事项和warning，大家可以自行查看文档。下面贴一段如何使用的示例代码，非常简单。

>>> import oneflow as flow
>>> import numpy as np

>>> input = flow.Tensor(np.arange(1, 5).reshape((1, 1, 4)), dtype=flow.float32)
>>> output = flow.nn.functional.interpolate(input, scale_factor=2.0, mode="linear")
>>> output
tensor([[[1.0000, 1.2500, 1.7500, 2.2500, 2.7500, 3.2500, 3.7500, 4.0000]]],
       dtype=oneflow.float32)

介绍完文档之后，我们看一下这个Op实现的Python前端接口，代码见：https://github.com/Oneflow-Inc/oneflow/blob/master/python/oneflow/nn/modules/interpolate.py#L25-L193 。这里的主要逻辑就是在根据是否传入了recompute_scale_factor参数来重新计算scale_factor的值，在获得了scale_factor之后根据传入的mode调用不同的插值Kernel的实现。见：

if len(x.shape) == 3 and self.mode == "nearest":
            return flow._C.upsample_nearest_1d(
                x, scale_factor=scale_factors[0], data_format="channels_first"
            )
        if len(x.shape) == 4 and self.mode == "nearest":
            return flow._C.upsample_nearest_2d(
                x,
                height_scale=scale_factors[0],
                width_scale=scale_factors[1],
                data_format="channels_first",
            )
        if len(x.shape) == 5 and self.mode == "nearest":
            return flow._C.upsample_nearest_3d(
                x,
                depth_scale=scale_factors[0],
                height_scale=scale_factors[1],
                width_scale=scale_factors[2],
                data_format="channels_first",
            )
        if len(x.shape) == 3 and self.mode == "area":
            assert output_size is not None
            return flow._C.adaptive_avg_pool1d(x, output_size)
        if len(x.shape) == 4 and self.mode == "area":
            assert output_size is not None
            return flow._C.adaptive_avg_pool2d(x, output_size)
        if len(x.shape) == 5 and self.mode == "area":
            assert output_size is not None
            return flow._C.adaptive_avg_pool3d(x, output_size)
        if len(x.shape) == 3 and self.mode == "linear":
            assert self.align_corners is not None
            return flow._C.upsample_linear_1d(
                x,
                scale_factor=scale_factors[0],
                align_corners=self.align_corners,
                data_format="channels_first",
            )
        if len(x.shape) == 4 and self.mode == "bilinear":
            assert self.align_corners is not None
            return flow._C.upsample_bilinear_2d(
                x,
                height_scale=scale_factors[0],
                width_scale=scale_factors[1],
                align_corners=self.align_corners,
                data_format="channels_first",
            )
        if len(x.shape) == 4 and self.mode == "bicubic":
            assert self.align_corners is not None
            return flow._C.upsample_bicubic_2d(
                x,
                height_scale=scale_factors[0],
                width_scale=scale_factors[1],
                align_corners=self.align_corners,
                data_format="channels_first",
            )
        if len(x.shape) == 5 and self.mode == "trilinear":
            assert self.align_corners is not None
            return flow._C.upsample_trilinear_3d(
                x,
                depth_scale=scale_factors[0],
                height_scale=scale_factors[1],
                width_scale=scale_factors[2],
                align_corners=self.align_corners,
                data_format="channels_first",
            )

所以Python前端就是处理了一些参数关系，然后调用了C++层的API来完成真正的计算过程。下面我们将分别介绍各种插值算法的原理以及在OneFlow中的实现。

0x2. AlignCorners解释

在上面的接口中，align_corners是一个非常重要的参数，这里我们先解释一下这个参数是什么含义再继续讲解每种Kernel的实现。这里以一张图片的nearest插值为例讲解align_corners的具体含义。

假设原始图像的大小是 $m\\times n$ ，目标图像是 $a\\times b$ ，那么两幅图像的边长比分别是 $m / a$ 和 $n / b$ 。那么目标图像的 $(i, j)$ 位置的像素可以通过上面的边长比对应回原图像，坐标为 $(i * m / a, j * n / b)$ 。当然这样获得的坐标可能不是整数，如果强行取整就是普通的最邻近插值，而双线性插值就是通过寻找距离这个对应坐标最近的四个像素点，来计算该点的值，如果坐标是 $(2.5, 4.5)$ ，那么最近的四个像素是 $(2 ， 4), (2 ， 5)$ , $(3 ， 4)$ ， $(3 ， 5)$ 。如果图形是灰度图，那么 $(i, j)$ 点的像素值可以通过下面的公式计算：
$f (i, j) = w 1 * p 1 + w 2 * p 2 + w 3 * p 3 + w 4 * p 4$
其中， $p i = (1, 2, 3, 4)$ 为最近的 $4$ 个像素点， $w_i$ 为各点的权重。

到这里并没有结束，我们需要特别注意的是，仅仅按照上面得到公式实现的双线性插值的结果和OpenCV/Matlab的结果是对应不起来的，这是为什么呢？

原因就是因为坐标系的选取问题，按照一些网上的公开实现，将源图像和目标图像的原点均选在左上角，然后根据插值公式计算目标图像每个点的像素，假设我们要将 $5\\times 5$ 的图像缩小成 $3\\times 3$ ，那么源图像和目标图像的对应关系如下图所示：

可以看到如果选择了左上角作为原点，那么最右边和最下边的像素是没有参与计算的，所以我们得到的结果和OpenCV/MatLab中的结果不会一致，那应该怎么做才是对的呢？

答案就是让两个图像的几何中心重合，并且目标图像的每个像素之间都是等间隔的，并且都和两边有一定的边距。如下图所示：

所以，我们只需要在计算坐标的时候将：

int x=i*m/a;
int y=j*n/b;

改成：

int x=(i+0.5)*m/a-0.5;
int y=(j+0.5)*n/b-0.5;

所以在interpolate Op的实现中提供了align_corners这个参数让用户选择是否对齐输入和输出的几何中心。

0x3. Linear插值

Linaer插值即线性插值。线性插值的几何意义即为概述图中利用过A点和B点的直线来近似表示原函数。如下图所示：

由于 $y-y_0)/(x-x_0)=(y_1-y_0)/(x_1-x_0)$ ，
那么 $x-x_0)/(x_1-x_0)=(y-y_0)/(y_1-y_0)=k$
再展开一下可得： $y=(1-k)*y_0+k*y_1$

在OneFlow中实现线性插值的代码在https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/upsample_linear_1d_kernel.cpp，我们只看前向，代码中的h1lambda就对应了这个公式里面的 $k$ 。

template<typename T>
OF_DEVICE_FUNC T GetLinearInputIndex(const int64_t out_dim_idx, const T scale, bool align_corners) {
  if (align_corners) {
    return static_cast<T>(scale * out_dim_idx);
  } else {
    T src_idx = scale * (out_dim_idx + 0.5) - 0.5;
    return static_cast<T>(src_idx < 0 ? 0 : src_idx);
  }
}

static void UpsampleLinear1DForward(const int64_t elem_cnt, const T* in_dptr,
                                    NdIndexOffsetHelper<int64_t, 3> in_helper,
                                    NdIndexOffsetHelper<int64_t, 3> out_helper, const int in_height,
                                    const float scale_factor, 以上是关于以OneFlow为例梳理深度学习框架的那些插值方法的主要内容，如果未能解决你的问题，请参考以下文章 
 OneFlow源码解析：算子签名的自动推断
 Autograd解析｜OneFlow学习笔记
 深度学习框架量化感知训练的思考及OneFlow的解决方案
 深度学习框架量化感知训练的思考及OneFlow的一种解决方案
 训练GPT-3，为什么原有的深度学习框架吃不消？
 一个算子在深度学习框架中的旅程