python - Pytorch 3D 张量与 1D 张量的内积生成 2D 张量

Question

操作:我有pytorch张量A尺寸[n x m x c]和B尺寸[1 x 1 x c] 。我想获取每个 1 x 1 x c 的内积矢量来自 A与 B从而生成一个张量 C尺寸[n x m] .

在我的网络的内部前向函数中，在特定步骤中我收到尺寸为 [N, channels, Height, Width] 的张量其中N是图像的数量，channels是特征图中的 channel 数，高度和宽度是当前特征图的高度和宽度。我还有一个[N x channels]来自其他子网络的特征图。下一步我要执行上述操作。

有人可以解释一下 pytorch 中实现这一步骤的最佳方法和功能吗？

我是 pytorch 新手，无法找到合适的方法。 Tensorflow支持NHWC格式，但我认为pytorch不支持，所以方法之一是将其 reshape 为[N, Height, Width, channels]然后迭代如下:

# if img is reshaped to [N, H, W, C]
img
# tensor of dimension [N, C]
aud

ans = torch.empty(N, H, W, dtype=torch.double)

for batches in range(img.shape[0]):
    for i in range(img.shape[1]):
        for j in range(img.shape[2]):
            ans[batches][i][j] = torch.dot(img[batches][i][j], aud[batches])

还有其他更干净的 API 吗？

PS:DeepMind 的论文“Object That Sound”中需要此步骤来进行声音定位步骤。

Answer 1

有一句台词

 ans = torch.einsum('nhwc,nc->nhw', img, aud)

torch.einsum的API如果您以前没有任何经验，可能很难掌握它，但它非常强大，并且概括了大量线性代数运算(转置、矩阵乘法和迹)。

import torch

N, H, W, C = 10, 11, 12, 13
img = torch.randn(N, H, W, C)
aud = torch.randn(N, C)

ans = torch.empty(N, H, W)
for batches in range(img.shape[0]):
    for i in range(img.shape[1]):
        for j in range(img.shape[2]):
            ans[batches][i][j] = torch.dot(img[batches][i][j], aud[batches])

ans2 = torch.einsum('nhwc,nc->nhw', img, aud)

assert torch.allclose(ans, ans2, atol=1e-6)

注意，由于数值精度问题，我必须将断言容差提高到标准 1e-8 以上。如果 einsum 成为更高级用例的瓶颈，请查看 opt_einsum它优化了底层操作的顺序以提高性能。