PyTorch lobpcg处理GPU复数稀疏数组报错的解决求助

问题场景 #

使用PyTorch的lobpcg函数在GPU上对角化大型稀疏数组时，处理实数数组正常，但处理复数数组时触发报错。相关代码片段如下：

from scipy import sparse
import torch
from torch import lobpcg
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def Torchfy(arr): # 将稀疏数组转移到GPU
    arr = arr.tocoo() # 转换为坐标格式
    arr = torch.sparse_coo_tensor(indices=torch.tensor([arr.row, arr.col]), values=torch.tensor(arr.data), size=arr.shape).to(device)
    return arr

def H_L(xcoor, ycoor, zcoor, b = 1, g1 = 1, g2 = 1, g3 = 1, e = 1, m = 1, hbar = 1):
    h2 = hbar ** 2
    c1 = - h2/(2*m) * (g1 + 2.5 * g2)
    c2 = g2 * h2 / m
    c3 = g3 * h2/(2*m)
    e2 = e ** 2
    eh = e / hbar
    e2h2 = e2 / h2

    y = 0.5 * b * sparse.spdiags(ycoor.reshape(N3), np.array([0]), N3, N3)
    x = 0.5 * b * sparse.spdiags(xcoor.reshape(N3), np.array([0]), N3, N3)
    y2 = y ** 2
    x2 = x ** 2

    lapl = ddxx + ddyy + ddzz
    a_nabl = -1j * eh * (x * ddy - y * ddx)
    a2 = e2h2 * (y2 + x2)
    tmp1 = c1 * (- lapl + a_nabl + a2)
    tmp1 = sparse.kron(tmp1, np.eye(4))

    kx2 = - ddxx + 1j * eh * y * ddx + e2h2 * y2
    ky2 = - ddyy - 1j * eh * x * ddy + e2h2 * x2
    tmp2 = sparse.kron(kx2, J_x2) + sparse.kron(ky2, J_y2) + sparse.kron(-ddzz, J_z2)

    kxky = 0.5 * (- 2 * ddxy - 2 * e2h2 * x * y - 1j * eh *( x * ddx - y * ddy))
    kxkz = 0.5 * (- 2 * ddxz + 1j * eh * y * ddz)
    kykz = 0.5 * (- 2 * ddyz - 1j * eh * x * ddz)
    tmp3 = sparse.kron(kxky, jxjy) + sparse.kron(kxkz, jxjz) + sparse.kron(kykz, jyjz)
    return tmp1 + tmp2 + tmp3

h = H_L(X, Y, Z)
h = Torchfy(h)
eigenvalues, eigenvectors = lobpcg(h, k=10, largest=False)

首次报错 #

RuntimeError: expected scalar type ComplexDouble but found Float

修改后二次报错 #

执行h = torch.complex(h.real, h.imag).to(device)后，触发新错误：

NotImplementedError: Could not run 'aten::view_as_real' with arguments from the 'SparseCPU' backend

解决方案 #

错误根源 #

1. 首次错误：转换为PyTorch稀疏张量时，默认采用Float类型，但复数矩阵需要ComplexDouble（torch.complex128）类型。
2. 二次错误：PyTorch对CPU后端的稀疏复数张量支持不完善，view_as_real操作未实现；且事后转换张量类型的方式会导致张量暂时回到CPU，引发后端不兼容问题。

修复步骤 #

直接在创建稀疏张量时指定复数类型并绑定GPU设备，修改Torchfy函数：

def Torchfy(arr): # 将稀疏数组转移到GPU
    arr = arr.tocoo() # 转换为坐标格式
    # 直接创建复数类型的稀疏张量并指定GPU设备
    arr = torch.sparse_coo_tensor(
        indices=torch.tensor([arr.row, arr.col], dtype=torch.long),
        values=torch.tensor(arr.data, dtype=torch.complex128),
        size=arr.shape,
        device=device
    )
    return arr

关键说明 #

• 显式指定values的dtype为torch.complex128，匹配复数矩阵的类型要求。
• 直接在torch.sparse_coo_tensor中指定device=device，避免事后调用.to(device)时可能的类型/后端冲突。
• 确保indices的类型为torch.long，符合PyTorch稀疏张量的索引类型规范。

修改后重新运行代码，lobpcg即可正常处理GPU上的稀疏复数矩阵。

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