Question

我正在尝试从3d张量向每个2d张量添加2d张量。假设我有一个a形状的张量(2,3,2)和一个b形状的张量(2,2)。

a = [[[1,2],
      [1,2],
      [1,2]],
     [[3,4],
      [3,4],
      [3,4]]]

b = [[1,2], [3,4]]

#the result i want to get
a[:, 0, :] + b
a[:, 1, :] + b
a[:, 2, :] + b

我想知道pytorch中是否有一种方法可以做到这一点。

Answer 1

最有效的方法是在b上添加额外的第二维，然后使用广播添加：

a = torch.Tensor([[[1,2],[1,2],[1,2]],[[3,4],[3,4],[3,4]]])
b = torch.Tensor([[1,2],[3,4]])
a += b.unsqueeze(1)

Answer 2

您想做什么：

a = [[[1,2],
      [1,2],
      [1,2]],
     [[3,4],
      [3,4],
      [3,4]]]

b = [[1,2], [3,4]]

a = torch.LongTensor(a)
b = torch.LongTensor(b)

a[:, 0, :] += b
a[:, 1, :] += b
a[:, 2, :] += b

print(a)

输出：

tensor([[[2, 4],
         [2, 4],
         [2, 4]],

        [[6, 8],
         [6, 8],
         [6, 8]]])

您可以执行以下相同操作。

a = (a.transpose(0, 1) + b).transpose(0, 1)
print(a) # prints the same tensor

Answer 3

@SinaAfrooze提出的解决方案是正确的，但它不是最快的。

TL; DR：torch.add(b.unsqueeze(1), a)更快。

基准：

import torch
a = torch.Tensor([[[1,2],[1,2],[1,2]],[[3,4],[3,4],[3,4]]])
b = torch.Tensor([[1,2],[3,4]])
z = a + b.unsqueeze(1)

%timeit k = torch.add(b.unsqueeze(1), a)
4.08 µs ± 25.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit z = a + b.unsqueeze(1)
4.14 µs ± 29 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

torch.equal(k, z)
True

如何从3d张量向每个2d张量添加2d张量

3 个答案: