在pytorch闪电中定义两个张量的以下乘法

Posted 2023-03-27

【中文标题】在pytorch闪电中定义两个张量的以下乘法

【英文标题】：Define following multiplication of two tensors in pytorch lightning

【发布时间】：2021-11-06 23:13:18

【问题描述】：

我想通过以下方式将以下两个张量 x（形状为 (BS,N,C)）和 y（形状为 (BS,1,C)）相乘：

BS = x.shape[0]
N = x.shape[1]
out = torch.zeros(size=x.shape)
for i in range(BS):
    for j in range(N):
        out[i, j, :] = torch.mul(x[i, j, :], y[i, 0, :])
return out

以这种方式实现它会产生错误“RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!（在检查方法 wrapper_native_layer_norm 中的参数权重时）”

通过设置修复时

out = torch.zeros(size=x.shape).to('cuda)

然后训练需要很长时间，因为我的 for 循环不是并行执行的。

所以我的问题是如何以 pytorch-lightning 的方式实现上面的两个 for 循环，这样我就可以定义函数 x = multiply_as_above(x, y) 并在我的神经网络的 feedword(self) 方法中使用它. 顺便说一句，上面定义的操作在我看来就像内核大小为 1 的卷积。也许我可以使用它？

【参考方案1】：

x*y 有什么问题吗？正如您在下面的代码中看到的，它产生的输出与您的函数完全相同：

import torch

torch.manual_seed(2021)

BS = 2
N = 3
C = 4

x = torch.rand(BS, N, C)
y = torch.rand(BS, 1, C)

# your function
def f(x, y):
    BS = x.shape[0]
    N = x.shape[1]
    out = torch.zeros(size=x.shape)
    for i in range(BS):
        for j in range(N):
            out[i, j, :] = torch.mul(x[i, j, :], y[i, 0, :])
    return out

out1 = f(x, y)
out2 = x*y

# comparing the outputs, we can see that they are identical
torch.all(out1 == out2)
# > tensor(True)

【讨论】：

谢谢，这很有帮助。你知道 out3=x * y.expand_as(x) 的区别吗，反向传播哪个更好？如果我没记错的话，使用 out3 的方法将 y 列复制 N 次。因此它需要更多的内存。

@mathemagier 不同之处在于您正在显式地进行扩展，而在 x*y 中它是作为广播隐式完成的。请考虑投票和/或标记为有帮助的答案。