双层循环结合梯度计算的PyTorch代码性能优化求助

我正在运行以下PyTorch代码：

for ie in range(100):
    energy = torch.tensor([.1 + .1 * ie], dtype = torch.float32, requires_grad = True)
    xe = .5 + 20 * (torch.log(energy) - log_01) / (log_10 - log_01) #(log_ parts are some constants)
    xe2 = xe * xe
    
    for it in range(100):
        theta = torch.tensor([it * theta_max / 99], dtype = torch.float32, requires_grad = True)
        Y[0] = torch.exp(PXmg1_p[0, 0]) + torch.exp(PXmg1_p[0, 1] * torch.pow(xe, PXmg1_p[0, 2]))
        Y[1] = torch.exp(PXmg2_p[0, 0]) + torch.exp(PXmg2_p[0, 1] * torch.pow(xe, PXmg2_p[0, 2]))
        Y[2] = torch.exp(PXmg3_p[0, 0]) + torch.exp(PXmg3_p[0, 1] * torch.pow(xe, PXmg3_p[0, 2]))
        Y[3] = torch.exp(PXmg4_p[0, 0]) + torch.exp(PXmg4_p[0, 1] * torch.pow(xe, PXmg4_p[0, 2]))

        thisp0_mg[ie, it] = solvecubic(energy, theta, 0)

        thisp0de_mg[ie, it] = solvecubic(energy, theta, 2)
        thisp0de2_mg[ie, it] = solvecubic(energy, theta, 22)
        thisp0de3_mg[ie, it] = solvecubic(energy, theta, 25)
        thisp0dth_mg[ie, it] = solvecubic(energy, theta, 3)
        thisp0dth2_mg[ie, it] = solvecubic(energy, theta, 32)

其中solvecubic函数定义如下：

def solvecubic(energy, theta, mode):
#Evaluate B given Y
B = torch.linalg.solve(A, Y)

val = 0
x = .5 + 4 * theta / theta_max
    
for i in range(4):
    val += B[i] * x**i

if mode == 0:
    return val.item()

elif mode == 2 or mode == 22 or mode == 25:
    if mode == 2:
        return torch.autograd.grad(val, energy, retain_graph = True)[0].item()

    if mode == 22:
        first_der = torch.autograd.grad(val, energy, create_graph = True)[0]

        return torch.autograd.grad(first_der, energy, retain_graph = True)[0].item()

    if mode == 25:
        first_der = torch.autograd.grad(val, energy, create_graph = True)[0]
        
        second_der = torch.autograd.grad(first_der, energy, create_graph = True)[0]

        return torch.autograd.grad(first_der, energy, retain_graph = True)[0].item()

elif mode == 3 or mode == 32:
    if mode == 3:
        return torch.autograd.grad(val, theta, retain_graph = True)[0].item()

    if mode == 32:
        first_der = torch.autograd.grad(val, theta, create_graph = True)[0]
        
        return torch.autograd.grad(first_der, theta, retain_graph = True)[0].item()

A为4x4范德蒙德矩阵。当前代码运行极慢，我的设备无CUDA环境，无法使用torch.cuda.clear_cache()，尝试使用del theta、del energy等操作后性能提升不明显。请问如何在每次迭代后删除计算图以节省内存、提升运行速度？

补充说明：内层循环次数设为1、2、3时，运行时间分别约为1分钟、2分30秒、10分钟。

1. 移除不必要的retain_graph=True #

你的代码中多次使用retain_graph=True，这会强制保留计算图导致内存堆积。由于每次迭代都是独立计算，完全不需要保留计算图，直接将所有retain_graph=True改为retain_graph=False（或直接省略，默认值即为False）。

2. 用torch.no_grad()包裹无需求导的代码块 #

对于仅需要数值输出的分支（如mode=0），用torch.no_grad()包裹，避免构建不必要的计算图：

if mode == 0:
    with torch.no_grad():
        return val.item()

3. 迭代后手动释放梯度并触发垃圾回收 #

在内层循环末尾添加代码，手动清空变量梯度并触发Python垃圾回收，及时释放内存：

# 清空梯度
energy.grad = None
theta.grad = None
# 触发垃圾回收
import gc
gc.collect()

4. 预计算常量，减少重复计算 #

xe是外层循环变量，内层循环中多次计算torch.pow(xe, ...)，可以在外层循环预计算这些幂次：

# 外层循环内添加
xe_pow1 = torch.pow(xe, PXmg1_p[0,2])
xe_pow2 = torch.pow(xe, PXmg2_p[0,2])
xe_pow3 = torch.pow(xe, PXmg3_p[0,2])
xe_pow4 = torch.pow(xe, PXmg4_p[0,2])

内层循环直接使用预计算结果：

Y[0] = torch.exp(PXmg1_p[0, 0]) + torch.exp(PXmg1_p[0, 1] * xe_pow1)
Y[1] = torch.exp(PXmg2_p[0, 0]) + torch.exp(PXmg2_p[0, 1] * xe_pow2)
Y[2] = torch.exp(PXmg3_p[0, 0]) + torch.exp(PXmg3_p[0, 1] * xe_pow3)
Y[3] = torch.exp(PXmg4_p[0, 0]) + torch.exp(PXmg4_p[0, 1] * xe_pow4)

5. 合并梯度计算，减少重复调用 #

当前内层循环调用6次solvecubic，每次都重复计算B和val，可以合并为一次计算，一次性求出所有需要的导数：

# 内层循环替换原有6次调用
theta = torch.tensor([it * theta_max / 99], dtype = torch.float32, requires_grad = True)
# 计算Y
Y[0] = torch.exp(PXmg1_p[0, 0]) + torch.exp(PXmg1_p[0, 1] * xe_pow1)
Y[1] = torch.exp(PXmg2_p[0, 0]) + torch.exp(PXmg2_p[0, 1] * xe_pow2)
Y[2] = torch.exp(PXmg3_p[0, 0]) + torch.exp(PXmg3_p[0, 1] * xe_pow3)
Y[3] = torch.exp(PXmg4_p[0, 0]) + torch.exp(PXmg4_p[0, 1] * xe_pow4)
# 一次性计算B和val
B = torch.linalg.solve(A, Y)
x = .5 + 4 * theta / theta_max
val = B[0] + B[1]*x + B[2]*x**2 + B[3]*x**3

# 计算所有需要的结果
thisp0_mg[ie, it] = val.item()
# energy的一阶导数
de1 = torch.autograd.grad(val, energy, create_graph=True)[0].item()
thisp0de_mg[ie, it] = de1
# energy的二阶导数
de2 = torch.autograd.grad(de1, energy, create_graph=True)[0].item()
thisp0de2_mg[ie, it] = de2
# energy的三阶导数（修正原mode=25的笔误）
de3 = torch.autograd.grad(de2, energy, retain_graph=False)[0].item()
thisp0de3_mg[ie, it] = de3
# theta的一阶导数
dth1 = torch.autograd.grad(val, theta, create_graph=True)[0].item()
thisp0dth_mg[ie, it] = dth1
# theta的二阶导数
dth2 = torch.autograd.grad(dth1, theta, retain_graph=False)[0].item()
thisp0dth2_mg[ie, it] = dth2

6. 修正原代码笔误 #

原solvecubic函数的mode=25分支中，错误地返回了一阶导数的梯度（实际是二阶导数），应该改为返回二阶导数的梯度（三阶导数），否则结果不符合预期。

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