如何在 Pytorch 中使用参数 batch_dims 实现 tf.gather_nd？

Question

我一直在做一个图像匹配项目，所以我需要找到 2 个图像之间的对应关系。要获得描述符，我需要一个插值函数。然而，当我读到一个在 Tensorflow 中完成的等效函数时，我仍然不知道如何在 Pytorch 中实现 tf.gather_nd(parmas, indices, barch_dims)。特别是当有参数时：batch_dims。我已经通过stackoverflow，还没有完美的对等。

Tensorflow 中引用的插值函数如下，我一直在尝试在 Pytorch Arguments 中实现它，信息如下：

inputs 是来自批量大小的 for 循环的密集特征图[i]，这意味着它是 3D[H, W, C]（在 pytorch 中是 [C, H, W]）

pos 是一组随机点坐标形状，如 [[i, j], [i, j],...,[i, j]]，因此在插值函数中（在 pytorch是 [[i,i,...,i], [j,j,...,j]])

然后当它们进入这个函数时，它会扩展它们的两个维度

我只想要一个带有参数 batch_dims 的 tf.gather_nd 的完美实现。谢谢！ 这是一个使用它的简单示例：

pos = tf.ones((12, 2)) ## stands for a set of coordinates [[i, i,…, i], [j, j,…, j]]
inputs = tf.ones((4, 4, 128)) ## stands for [H, W, C] of dense feature map
outputs = interpolate(pos, inputs, batched=False)
print(outputs.get_shape()) # We get (12, 128) here

插值函数（tf 版本）：

def interpolate(pos, inputs, nd=True):

    pos = tf.expand_dims(pos, 0)
    inputs = tf.expand_dims(inputs, 0)

    h = tf.shape(inputs)[1]
    w = tf.shape(inputs)[2]

    i = pos[:, :, 0]
    j = pos[:, :, 1]

    i_top_left = tf.clip_by_value(tf.cast(tf.math.floor(i), tf.int32), 0, h - 1)
    j_top_left = tf.clip_by_value(tf.cast(tf.math.floor(j), tf.int32), 0, w - 1)

    i_top_right = tf.clip_by_value(tf.cast(tf.math.floor(i), tf.int32), 0, h - 1)
    j_top_right = tf.clip_by_value(tf.cast(tf.math.ceil(j), tf.int32), 0, w - 1)

    i_bottom_left = tf.clip_by_value(tf.cast(tf.math.ceil(i), tf.int32), 0, h - 1)
    j_bottom_left = tf.clip_by_value(tf.cast(tf.math.floor(j), tf.int32), 0, w - 1)

    i_bottom_right = tf.clip_by_value(tf.cast(tf.math.ceil(i), tf.int32), 0, h - 1)
    j_bottom_right = tf.clip_by_value(tf.cast(tf.math.ceil(j), tf.int32), 0, w - 1)

    dist_i_top_left = i - tf.cast(i_top_left, tf.float32)
    dist_j_top_left = j - tf.cast(j_top_left, tf.float32)
    w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
    w_top_right = (1 - dist_i_top_left) * dist_j_top_left
    w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
    w_bottom_right = dist_i_top_left * dist_j_top_left

    if nd:
        w_top_left = w_top_left[..., None]
        w_top_right = w_top_right[..., None]
        w_bottom_left = w_bottom_left[..., None]
        w_bottom_right = w_bottom_right[..., None]

    interpolated_val = (
        w_top_left * tf.gather_nd(inputs, tf.stack([i_top_left, j_top_left], axis=-1), batch_dims=1) +
        w_top_right * tf.gather_nd(inputs, tf.stack([i_top_right, j_top_right], axis=-1), batch_dims=1) +
        w_bottom_left * tf.gather_nd(inputs, tf.stack([i_bottom_left, j_bottom_left], axis=-1), batch_dims=1) +
        w_bottom_right * tf.gather_nd(inputs, tf.stack([i_bottom_right, j_bottom_right], axis=-1), batch_dims=1)
    )

    interpolated_val = tf.squeeze(interpolated_val, axis=0)
    return interpolated_val

Answer 1

据我所知，在 PyTorch 中没有直接等同于 tf.gather_nd 的版本，使用 batch_dims 实现通用版本并不是那么简单。但是，您可能不需要通用版本，并且考虑到您的 interpolate 函数的上下文，[C, H, W] 的版本就足够了。

在interpolate的开头你在前面添加了一个奇异维度，也就是批量维度。在tf.gather_nd 中设置batch_dims=1 意味着在开始时有一个批次维度，因此它每批次都应用它，即它索引inputs[0] 和pos[0] 等。添加单个批次维度没有任何好处，因为您可以直接使用直接计算。

# Adding singular batch dimension
# Shape: [1, num_pos, 2]
pos = tf.expand_dims(pos, 0)
# Shape: [1, H, W, C]
inputs = tf.expand_dims(inputs, 0)


batched_result = tf.gather_nd(inputs, pos, batch_dims=1)
single_result = tf.gater_nd(inputs[0], pos[0])

# The first element in the batched result is the same as the single result
# Hence there is no benefit to adding a singular batch dimension.
tf.reduce_all(batched_result[0] == single_result) # => True

单机版

在 PyTorch 中，[H, W, C] 的实现可以通过 Python 的索引来完成。虽然 PyTorch 通常使用 [C, H, W] 来存储图像，但这只是索引的维度的问题，但为了比较，我们让它们与 TensorFlow 中的相同。如果您要手动索引它们，您可以这样做：inputs[pos_h[0], pos_w[0]]、inputs[pos_h[1], pos_w[1]] 等等。 PyTorch 允许您通过将索引提供为列表来自动执行此操作：inputs[pos_h, pos_w]，其中pos_h 和pos_w 具有相同的长度。您需要做的就是将您的 pos 拆分为两个单独的张量，一个用于沿高度维度的索引，另一个用于沿宽度维度的索引，您在 TensorFlow 版本中也这样做了。

inputs = torch.randn(4, 4, 128)
# Random positions 0-3, shape: [12, 2]
pos = torch.randint(4, (12, 2))

# Positions split by dimension
pos_h = pos[:, 0]
pos_w = pos[:, 1]

# Index the inputs with the indices per dimension
gathered = inputs[pos_h, pos_w]

# Verify that it's identical to TensorFlow's output
inputs_tf = tf.convert_to_tensor(inputs.numpy())
pos_tf = tf.convert_to_tensor(pos.numpy())
gathered_tf = tf.gather_nd(inputs_tf, pos_tf)
gathered_tf = torch.from_numpy(gathered_tf.numpy())

torch.equal(gathered_tf, gathered) # => True

如果您想将其应用于大小为 [C, H, W] 的张量，则只需更改要索引的维度：

# For [H, W, C]
gathered = inputs[pos_h, pos_w]

# For [C, H, W]
gathered = inputs[:, pos_h, pos_w]

批处理版

使其成为批处理版本（用于[N, H, W, C] 或[N, C, H, W]）并不难，使用它更合适，因为无论如何您都在处理批处理。唯一棘手的部分是批次中的每个元素都应该只应用于相应的批次。为此，需要枚举批次维度，这可以使用torch.arange 完成。批处理枚举只是带有批处理索引的列表，它将与pos_h 和pos_w 索引组合，产生inputs[0, pos_h[0, 0], pos_h[0, 0]]、inputs[0, pos_h[0, 1], pos_h[0, 1]] ... inputs[1, pos_h[1, 0], pos_h[1, 0]] 等。

batch_size = 3
inputs = torch.randn(batch_size, 4, 4, 128)
# Random positions 0-3, different for each batch, shape: [3, 12, 2]
pos = torch.randint(4, (batch_size, 12, 2))

# Positions split by dimension
pos_h = pos[:, :, 0]
pos_w = pos[:, :, 1]

batch_enumeration = torch.arange(batch_size) # => [0, 1, 2]
# pos_h and pos_w have shape [3, 12], so the batch enumeration needs to be
# repeated 12 times per batch.
# Unsqueeze to get shape [3, 1], now the 1 could be repeated to 12, but
# broadcasting will do that automatically.
batch_enumeration = batch_enumeration.unsqueeze(1)
# Index the inputs with the indices per dimension
gathered = inputs[batch_enumeration, pos_h, pos_w]

# Again, verify that it's identical to TensorFlow's output
inputs_tf = tf.convert_to_tensor(inputs.numpy())
pos_tf = tf.convert_to_tensor(pos.numpy())
# This time with batch_dims=1
gathered_tf = tf.gather_nd(inputs_tf, pos_tf, batch_dims=1)
gathered_tf = torch.from_numpy(gathered_tf.numpy())

torch.equal(gathered_tf, gathered) # => True

同样，对于[N, C, H, W]，只需要更改索引的维度：

# For [N, H, W, C]
gathered = inputs[batch_enumeration, pos_h, pos_w]

# For [N, C, H, W]
gathered = inputs[batch_enumeration, :, pos_h, pos_w]

---

只是对interpolate 实现的一点注解，四舍五入位置（分别为 floor 和 ceil ）没有意义，因为索引必须是整数，所以它没有效果，只要你的位置是实际的索引.这也导致 i_top_left 和 i_bottom_left 是相同的值，但即使它们要以不同的方式四舍五入，它们也总是相隔 1 个位置。此外，i_top_left 和 i_top_right 在字面上是相同的。我认为这个函数不会产生有意义的输出。我不知道您想要实现什么，但如果您正在寻找图像插值，您可以查看torch.nn.functional.interpolate。

Answer 2

当pos 是二维数组而不是一维数组（不包括批处理维度）时，这只是Michael Jungo 的批处理版本answer 的扩展。

bs = 2
H = 4
W = 6
C = 3
inputs = torch.randn(bs, H, W, C)
pos_h = torch.randint(H, (bs, H, W))
pos_w = torch.randint(W, (bs, H, W))
batch_enumeration = torch.arange(bs)
batch_enumeration = batch_enumeration.unsqueeze(1).unsqueeze(2)
inputs.shape
Out[34]: torch.Size([2, 4, 6, 3])
pos_h.shape
Out[35]: torch.Size([2, 4, 6])
pos_w.shape
Out[36]: torch.Size([2, 4, 6])
batch_enumeration.shape
Out[37]: torch.Size([2, 1, 1])
gathered = inputs[batch_enumeration, pos_h, pos_w]

频道优先，我们还需要枚举频道

inputs = torch.randn(bs, C, H, W)
pos_h = torch.randint(H, (bs, 1, H, W))
pos_w = torch.randint(W, (bs, 1, H, W))
batch_enumeration = torch.arange(bs)
batch_enumeration = batch_enumeration.unsqueeze(1).unsqueeze(2).unsqueeze(3)
channel_enumeration = torch.arange(C)
channel_enumeration = channel_enumeration.unsqueeze(0).unsqueeze(2).unsqueeze(3)
inputs.shape
Out[49]: torch.Size([2, 3, 4, 6])
pos_h.shape
Out[50]: torch.Size([2, 1, 4, 6])
pos_w.shape
Out[51]: torch.Size([2, 1, 4, 6])
batch_enumeration.shape
Out[52]: torch.Size([2, 1, 1, 1])
channel_enumeration.shape
Out[57]: torch.Size([1, 3, 1, 1])
gathered = inputs[batch_enumeration, channel_enumeration, pos_h, pos_w]
gathered.shape
Out[59]: torch.Size([2, 3, 4, 6])

让我们验证一下

inputs_np = inputs.numpy()
pos_h_np = pos_h.numpy()
pos_w_np = pos_w.numpy()
gathered_np = gathered.numpy()
pos_h_np[0,0,0,0]
Out[68]: 0
pos_w_np[0,0,0,0]
Out[69]: 3
inputs_np[0,:,0,3]
Out[71]: array([ 0.79122806, -2.190181  , -0.16741803], dtype=float32)
gathered_np[0,:,0,0]
Out[72]: array([ 0.79122806, -2.190181  , -0.16741803], dtype=float32)
pos_h_np[1,0,3,4]
Out[73]: 1
pos_w_np[1,0,3,4]
Out[74]: 2
inputs_np[1,:,1,2]
Out[75]: array([ 0.9282498 , -0.34945545,  0.9136222 ], dtype=float32)
gathered_np[1,:,3,4]
Out[77]: array([ 0.9282498 , -0.34945545,  0.9136222 ], dtype=float32)