评价此页

复数#

创建于:2025 年 6 月 16 日 | 最后更新于:2025 年 6 月 16 日

复数是可以表示为 a+bja + bj 形式的数字,其中 a 和 b 是实数,j 是虚数单位,满足方程 j2=1j^2 = -1。复数在数学和工程领域经常出现,尤其是在信号处理等主题中。传统上,许多用户和库(例如 TorchAudio)通过使用形状为 (...,2)(..., 2) 的浮点张量来表示复数,其中最后一个维度包含实部和虚部。

复数数据类型的张量在处理复数时提供了更自然的编程体验。复数张量上的操作(例如 torch.mv()torch.matmul())可能比模拟它们的浮点张量上的操作更快、内存效率更高。PyTorch 中涉及复数的操作经过优化,可以使用矢量化汇编指令和专用内核(例如 LAPACK,cuBlas)。

注意

torch.fft 模块中的频谱运算支持原生复数张量。

警告

复数张量是一项测试版功能,可能会发生更改。

创建复数张量#

我们支持两种复数数据类型:torch.cfloattorch.cdouble

>>> x = torch.randn(2,2, dtype=torch.cfloat)
>>> x
tensor([[-0.4621-0.0303j, -0.2438-0.5874j],
     [ 0.7706+0.1421j,  1.2110+0.1918j]])

注意

复数张量的默认数据类型由默认浮点数据类型决定。如果默认浮点数据类型是 torch.float64,则复数将被推断为 torch.complex128 数据类型,否则假定它们的数据类型为 torch.complex64

torch.linspace()torch.logspace()torch.arange() 外,所有工厂函数均支持复数张量。

从旧表示过渡#

对于那些目前通过形状为 (...,2)(..., 2) 的实数张量来处理缺乏复数张量问题的用户,可以使用 torch.view_as_complex()torch.view_as_real() 轻松地将其代码切换为使用复数张量。请注意,这些函数不会执行任何复制,而是返回输入张量的视图。

>>> x = torch.randn(3, 2)
>>> x
tensor([[ 0.6125, -0.1681],
     [-0.3773,  1.3487],
     [-0.0861, -0.7981]])
>>> y = torch.view_as_complex(x)
>>> y
tensor([ 0.6125-0.1681j, -0.3773+1.3487j, -0.0861-0.7981j])
>>> torch.view_as_real(y)
tensor([[ 0.6125, -0.1681],
     [-0.3773,  1.3487],
     [-0.0861, -0.7981]])

访问实部和虚部#

可以使用 realimag 属性访问复数张量的实部和虚部。

注意

访问 realimag 属性不会分配任何内存,并且对 realimag 张量的原地更新将修改原始复数张量。此外,返回的 realimag 张量不是连续的。

>>> y.real
tensor([ 0.6125, -0.3773, -0.0861])
>>> y.imag
tensor([-0.1681,  1.3487, -0.7981])

>>> y.real.mul_(2)
tensor([ 1.2250, -0.7546, -0.1722])
>>> y
tensor([ 1.2250-0.1681j, -0.7546+1.3487j, -0.1722-0.7981j])
>>> y.real.stride()
(2,)

角度和模#

可以使用 torch.angle()torch.abs() 计算复数张量的角度和模。

>>> x1=torch.tensor([3j, 4+4j])
>>> x1.abs()
tensor([3.0000, 5.6569])
>>> x1.angle()
tensor([1.5708, 0.7854])

线性代数#

许多线性代数运算,如 torch.matmul()torch.linalg.svd()torch.linalg.solve() 等,都支持复数。如果您想请求我们目前不支持的操作,请搜索是否已有相关 issue,如果没有,请提交一个

序列化#

复数张量可以被序列化,允许将数据保存为复数值。

>>> torch.save(y, 'complex_tensor.pt')
>>> torch.load('complex_tensor.pt')
tensor([ 0.6125-0.1681j, -0.3773+1.3487j, -0.0861-0.7981j])

自动微分#

PyTorch 支持复数张量的自动微分。计算出的梯度是共轭维丁格(Conjugate Wirtinger)导数,其负值正是梯度下降算法中使用的最陡下降方向。因此,所有现有的优化器都可以直接用于复数参数。更多详情,请参阅笔记 复数自动微分

优化器#

从语义上讲,我们将通过 PyTorch 优化器处理复数参数定义为等同于在复数参数的 torch.view_as_real() 等效表示上使用相同的优化器进行更新。更具体地说:

>>> params = [torch.rand(2, 3, dtype=torch.complex64) for _ in range(5)]
>>> real_params = [torch.view_as_real(p) for p in params]

>>> complex_optim = torch.optim.AdamW(params)
>>> real_optim = torch.optim.AdamW(real_params)

real_optimcomplex_optim 将在参数上计算相同的更新,尽管两者之间可能存在轻微的数值差异,类似于 foreach 与 forloop 优化器以及可捕获与默认优化器之间的数值差异。更多详情,请参阅 数值精度

具体来说,虽然您可以将我们优化器处理复数张量的方式视为分别优化它们的 p.realp.imag 部分,但实现细节并非完全如此。请注意,torch.view_as_real() 等效操作会将复数张量转换为一个形状为 (...,2)(..., 2) 的实数张量,而将复数张量拆分为两个张量则会得到两个大小为 (...)(...) 的张量。这一区别对逐元素优化器(如 AdamW)没有影响,但对于执行全局约简的优化器(如 LBFGS)可能会导致轻微的差异。我们目前没有执行每张量约简的优化器,因此尚未定义此行为。如果您有需要精确定义此行为的用例,请提交一个 issue。

我们尚未完全支持以下子系统:

  • 量化

  • JIT

  • 稀疏张量

  • 分布式

如果其中任何一项有助于您的用例,请搜索是否已有相关 issue,如果没有,请提交一个