Tensorflow / Numpy 中的 torch.nn.functional.grid_sample 等价物是啥？

Posted 2023-03-12

【中文标题】Tensorflow / Numpy 中的 torch.nn.functional.grid_sample 等价物是啥？

【英文标题】：What is the equivalent of torch.nn.functional.grid_sample in Tensorflow / Numpy?Tensorflow / Numpy 中的 torch.nn.functional.grid_sample 等价物是什么？

【发布时间】：2019-03-24 02:19:12

【问题描述】：

我是 pytorch 的新手，一直在尝试转换一些代码。找不到这个特定的功能。它存在于tensorflow中吗？

【问题讨论】：

不确定这是否可行，但您需要针对 2D 和 3D 案例的解决方案还是只需要其中一个案例？你还需要填充功能吗？

我在 tensorflow 中实现了interpn()，基本上就是这样做的。我希望这对您有用：github.com/adalca/neuron/blob/master/neuron/utils.py。另请参阅 layers.py 中的 transform() 和 SpatialTransformer，这可能与您有关。

【参考方案1】：

我认为 TensorFlow 中没有提供类似的东西。这是 2D 情况的可能实现（我没有考虑填充，但代码的行为应该类似于 border 模式）。请注意，与 PyTorch 版本不同，我假设输入维度顺序为 (batch_size, height, width, channels)（这在 TensorFlow 中很常见）。

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

def grid_sample_2d(inp, grid):
    in_shape = tf.shape(inp)
    in_h = in_shape[1]
    in_w = in_shape[2]

    # Find interpolation sides
    i, j = grid[..., 0], grid[..., 1]
    i = tf.cast(in_h - 1, grid.dtype) * (i + 1) / 2
    j = tf.cast(in_w - 1, grid.dtype) * (j + 1) / 2
    i_1 = tf.maximum(tf.cast(tf.floor(i), tf.int32), 0)
    i_2 = tf.minimum(i_1 + 1, in_h - 1)
    j_1 = tf.maximum(tf.cast(tf.floor(j), tf.int32), 0)
    j_2 = tf.minimum(j_1 + 1, in_w - 1)

    # Gather pixel values
    n_idx = tf.tile(tf.range(in_shape[0])[:, tf.newaxis, tf.newaxis], tf.concat([[1], tf.shape(i)[1:]], axis=0))
    q_11 = tf.gather_nd(inp, tf.stack([n_idx, i_1, j_1], axis=-1))
    q_12 = tf.gather_nd(inp, tf.stack([n_idx, i_1, j_2], axis=-1))
    q_21 = tf.gather_nd(inp, tf.stack([n_idx, i_2, j_1], axis=-1))
    q_22 = tf.gather_nd(inp, tf.stack([n_idx, i_2, j_2], axis=-1))

    # Interpolation coefficients
    di = tf.cast(i, inp.dtype) - tf.cast(i_1, inp.dtype)
    di = tf.expand_dims(di, -1)
    dj = tf.cast(j, inp.dtype) - tf.cast(j_1, inp.dtype)
    dj = tf.expand_dims(dj, -1)

    # Compute interpolations
    q_i1 = q_11 * (1 - di) + q_21 * di
    q_i2 = q_12 * (1 - di) + q_22 * di
    q_ij = q_i1 * (1 - dj) + q_i2 * dj

    return q_ij

# Test it
inp = tf.placeholder(tf.float32, [None, None, None, None])
grid = tf.placeholder(tf.float32, [None, None, None, 2])
res = grid_sample_2d(inp, grid)
with tf.Session() as sess:
    # Make test image
    im_grid_i, im_grid_j = np.meshgrid(np.arange(6), np.arange(10), indexing='ij')
    im = im_grid_i + im_grid_j
    im = im / im.max()
    im = np.stack([im] * 3, axis=-1)
    # Test grid 1: complete image
    grid1 = np.stack(np.meshgrid(np.linspace(-1, 1, 15), np.linspace(-1, 1, 18), indexing='ij'), axis=-1)
    # Test grid 2: lower right corner
    grid2 = np.stack(np.meshgrid(np.linspace(0, 1, 15), np.linspace(.5, 1, 18), indexing='ij'), axis=-1)
    # Run
    res1, res2 = sess.run(res, feed_dict=inp: [im, im], grid: [grid1, grid2])
    # Plot image and sampled grids
    plt.figure()
    plt.imshow(im)
    plt.figure()
    plt.imshow(res1)
    plt.figure()
    plt.imshow(res2)

这是生成的图像，首先是输入：

第一个网格结果，即第一个图像，但形状不同：

第二个网格结果，跨越右下角的一个区域：

【讨论】：

您的解决方案对我不起作用。但this 工作。我认为这是因为最后的插值部分是错误的

我赞成上述评论，因为上面提供的链接中的代码对我有用。我没有检查答案中提供的代码。请不要将我对上述评论的赞成误认为“答案中的代码不起作用”。

@Khan 不太清楚 bilinear_sampler 的输入是什么。

好的，@Khan 的建议函数似乎提供了与 numpy 的 grid_sample 相同的结果。对于 x 和 y，tf 版本的输入是 (b, h, w)。