PyTorch中CUDA张量负数开根号转ComplexFloatTensor实现方法

无需将张量转移到CPU，直接利用PyTorch对复数张量的原生支持即可完成计算，全程在GPU上运行：

• 先将包含正负值的CUDA实数张量转换为复数张量（虚部设为0）
• 调用torch.sqrt()计算复数平方根，结果自动为ComplexFloatTensor类型

代码示例 #

import torch

# 示例CUDA张量，包含正负数值
x = torch.tensor([-4.0, 9.0, -1.0, 16.0], device='cuda')

# 转换为复数张量（实部为原张量值，虚部为0）
x_complex = torch.complex(x, torch.zeros_like(x))

# 计算复数平方根，结果保留在CUDA设备上
sqrt_result = torch.sqrt(x_complex)

print(sqrt_result)
# 输出：tensor([0.+2.j, 3.+0.j, 0.+1.j, 4.+0.j], device='cuda:0')

简化写法（PyTorch 1.7+） #

可以用torch.as_complex()简化张量转换步骤，需要先将原张量拼接出虚部维度：

# 拼接实部和虚部（虚部为0），再转为复数张量
x_complex = torch.as_complex(torch.stack([x, torch.zeros_like(x)], dim=-1))
sqrt_result = torch.sqrt(x_complex)

注意：如果直接对含负数的实数张量调用torch.sqrt()会触发报错，因此必须先转为复数张量再计算。

内容的提问来源于stack exchange，提问作者Cantor