python - 如何在 PyTorch 中生成具有不同向量的新张量？

Question

我想用 a 和 b 生成新的 a○b 向量(○ 表示按元素相乘)。我的代码如下，但由于 for 的原因，性能看起来很糟糕。有什么有效的方法吗？

a = torch.rand(batch_size, a_len, hid_dim)
b = torch.rand(batch_size, b_len, hid_dim)
# a_elmwise_mul_b = torch.zeros(batch_size, a_len, b_len, hid_dim)
for sample in range(batch_size):
    for ai in range(a_len):
        for bi in range(b_len):
            a_elmwise_mul_b[sample, ai, bi] = torch.mul(a[sample, ai], b[sample, bi])

更新

我更新了我的代码，引用艾哈迈德!谢谢。

N = 16
hid_dim = 50
a_seq_len = 10
b_seq_len = 20
a = torch.randn(N, a_seq_len, hid_dim)
b = torch.randn(N, b_seq_len, hid_dim)
shape = (N, a_seq_len, b_seq_len, hid_dim)

a_dash = a.unsqueeze(2) # (N, a_len, 1,     hid_dim)
b_dash = b.unsqueeze(1) # (N, 1,     b_len, hid_dim)
a_dash = a_dash.expand(shape)
b_dash = b_dash.expand(shape)
print(a_dash.size(), b_dash.size())
mul = a_dash * b_dash
print(mul.size())
----------
torch.Size([16, 10, 20, 50]) torch.Size([16, 10, 20, 50])
torch.Size([16, 10, 20, 50])

Answer 1

根据您的问题定义，您似乎想要将两个张量相乘，例如形状为 AxE 和 的 A 和 B BxE 并想要获得形状为 AxBxE 的张量。这意味着您想要将张量 A 的每一行与整个张量 B 相乘。如果它是正确的，那么我们就不称其为逐元素乘法。

您可以按如下方式实现您的目标。

import torch

# batch_size = 16, a_len = 10, b_len = 20, hid_dim = 50
a = torch.rand(16, 10, 50)
b = torch.rand(16, 20, 50)

c = a.unsqueeze(2).expand(*a.size()[:-1], b.size(1), a.size()[-1])
d = b.unsqueeze(1).expand(b.size()[0], a.size(1), *b.size()[1:])
print(c.size(), d.size())
print(c.size(), d.size())

mul = c * d       # shape of c, d: 16 x 10 x 20 x 50
print(mul.size()) # 16 x 10 x 20 x 50

这里，mul 张量是您想要的结果。澄清一下，上面两行涉及 c 和 d 计算，相当于:

c = a.unsqueeze(2).expand(a.size(0), a.size(1), b.size(1), a.size(2))
d = b.unsqueeze(1).expand(b.size(0), a.size(1), b.size(1), b.size(2))