pytorch - PyTorch 中复数的矩阵乘法

Question

我正在尝试在 PyTorch 中将两个复数矩阵相乘，看起来 the torch.matmul functions is not added yet to PyTorch library for complex numbers.

您有什么建议或有其他方法可以在 PyTorch 中乘以复杂的矩阵吗？

Answer 1

目前 torch.matmul 不支持复杂的张量，例如 ComplexFloatTensor，但您可以像下面的代码一样紧凑:

def matmul_complex(t1,t2):
    return torch.view_as_complex(torch.stack((t1.real @ t2.real - t1.imag @ t2.imag, t1.real @ t2.imag + t1.imag @ t2.real),dim=2))

尽可能避免使用 for 循环，因为这会导致执行速度变慢。矢量化是通过使用内置方法实现的，如我所附的代码所示。例如，对于 2 个尺寸为 1000 X 1000 的随机复杂矩阵，您的代码在 CPU 上大约需要 6.1 秒，而矢量化版本仅需要 101 毫秒(快约 60 倍)。

更新:

从 PyTorch 1.7.0 开始(如@EduardoReis 所述)，您可以在复数矩阵之间进行矩阵乘法，类似于实值矩阵，如下所示:

t1 @ t2(对于 t1、t2 复杂矩阵)。