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

Question

操作：我有尺寸[nxmxc]火炬张量A和尺寸[1 x 1 xc]火炬张量B 我想从A到B取每个1 x 1 xc向量的内积，并因此生成尺寸为[nxm]的张量C

在特定步骤的网络内部转发函数中，我收到张量为[N, channels, Height, Width]张量，其中N是图像数量，channels是特征图中的通道数，而height和width是当前特征的地图。 我也有其他子网的[N x channels]功能图。 在下一步中，我要执行上述操作。

有人可以解释pytorch中实现此步骤的最佳方法和功能。

我是pytorch的新手，无法找到正确的方法。 Tensorflow支持NHWC格式，但我认为pytorch不支持，因此一种方法是将其重塑为[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在更高级的用einsum成为瓶颈，请查看opt_einsum ，它可以优化底层操作的顺序以提高性能。