如何在 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])

更新

我参考 Ahmad 更新了我的代码！谢谢。

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 和 BxE 的 A 和 B，并希望得到形状为 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))

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

How to generate a new Tensor with different vectors in PyTorch?

python

pytorch

更新