在 Pytorch 中连接两个不同维度的张量

Question

是否可以在不使用 for 循环的情况下连接两个不同维度的张量。

例如张量 1 的尺寸为 (15, 200, 2048)，张量 2 的尺寸为 (1, 200, 2048)。是否可以沿第一张量中第一维的所有 15 个索引将第二张量与第一张量连接（沿张量 1 的第一维广播第二张量，同时沿第一张量的第三维连接）？生成的张量应该有尺寸 (15, 200, 4096)。

没有for循环是否可以做到这一点？

Answer 1

您可以在连接之前手动进行广播（使用Tensor.expand()）（使用torch.cat()）：

import torch

a = torch.randn(15, 200, 2048)
b = torch.randn(1, 200, 2048)

repeat_vals = [a.shape[0] // b.shape[0]] + [-1] * (len(b.shape) - 1)
# or directly repeat_vals = (15, -1, -1) or (15, 200, 2048) if shapes are known and fixed...
res = torch.cat((a, b.expand(*repeat_vals)), dim=-1)
print(res.shape)
# torch.Size([15, 200, 4096])