分批划分 3D 张量 (Pytorch)

Question

我有一个大小为 100x5x2 的 3D 张量和轴 = 1 上的张量的平均值，其形状为 100x2。 100 这里是批量大小。通常没有批处理，形状 5x2 和 2 的张量的划分工作完美，但在带有批处理的 3D 张量的情况下，我收到错误。

a = torch.rand(5,2)
b = torch.rand(2)
z=a/b

给了我预期的答案。

a = torch.rand(100,5,2)
b = torch.rand(100,2)
z=a/b

给我以下错误。

张量 a (5) 的大小必须与非单维 1 的张量 b (100) 的大小相匹配。

如何划分这些张量以使我的输出形状为 100x5x2 ？像 bmm 这样的除法？

Answer 1

简单地做：

z = a / b.unsqueeze(1)

这会在b 中添加一个额外的维度，并使其形状为 (100, 1, 2)，这与 a 的广播兼容。