自定义对比损失PyTorch反向传播过慢，求性能优化方法

我正在为网络实现自定义对比损失，但训练过程极慢，排查后发现自定义损失的反向传播是主要性能瓶颈。以下是简化后的复现代码：

import torch
import time
import random

n = 2000
neighbors = 2
x = torch.randn((n, n), requires_grad=True)

start = time.time()
row_sum = torch.sum(torch.exp(x), dim=1) - torch.exp(torch.diagonal(x, 0))
y = torch.zeros((n, neighbors))
for i in range(n):
    for j in range(neighbors):
        y[i][j] = x[i][random.randrange(0, n)]
loss = torch.sum(-torch.log(torch.exp(y) / row_sum.view(-1, 1)))  # Using broadcasting
print('Forward time: {:.3f}'.format(time.time() - start))

start = time.time()
loss.backward()
print('Backward time: {:.3f}'.format(time.time() - start))

测试结果 #

默认参数（n=2000，neighbors=2）下的运行耗时：

Forward time: 0.255
Backward time: 37.502

反向传播耗时与n（样本数）和neighbors（每个样本的邻居数）强相关，不同参数组合的测试结果如下：

• 当n = 1000、neighbors = 2时：

Forward time: 0.241
Backward time: 9.094

• 当n = 3000、neighbors = 2时：

Forward time: 0.250
Backward time: 136.768

• 当n = 2000、neighbors = 1时：

Forward time: 0.124
Backward time: 22.925

• 当n = 2000、neighbors = 3时：

Forward time: 0.210
Backward time: 58.341

性能优化方案 #

1. 向量化采样过程，消除嵌套循环 #

原代码的嵌套循环逐个赋值会破坏PyTorch计算图的优化逻辑，导致反向传播效率低下。可以通过生成批量索引，用张量的gather方法一次性提取所需元素：

import torch
import time
import random

n = 2000
neighbors = 2
x = torch.randn((n, n), requires_grad=True)

start = time.time()
row_sum = torch.sum(torch.exp(x), dim=1) - torch.exp(torch.diagonal(x, 0))
# 生成所有采样的列索引，转为张量
indices = torch.tensor([[random.randrange(0, n) for _ in range(neighbors)] for _ in range(n)])
# 批量提取x中对应位置的元素，避免循环
y = x.gather(1, indices)
loss = torch.sum(-torch.log(torch.exp(y) / row_sum.view(-1, 1)))
print('Forward time: {:.3f}'.format(time.time() - start))

start = time.time()
loss.backward()
print('Backward time: {:.3f}'.format(time.time() - start))

2. 简化损失计算，消除冗余操作 #

利用对数运算性质log(a/b) = log(a) - log(b)，可以去掉冗余的exp计算，直接简化损失公式：

# 简化后的损失计算，替换原loss行
loss = torch.sum(-(y - torch.log(row_sum).view(-1, 1)))

这不仅能加快前向传播，还能让反向传播的梯度计算逻辑更简洁，进一步降低耗时。

3. 优先使用PyTorch内置优化损失函数 #

如果你的损失逻辑符合标准对比学习范式，直接使用PyTorch内置的torch.nn.CrossEntropyLoss或torch.nn.NTCLoss等函数，这些底层实现经过高度优化，性能远优于自定义循环实现。

4. 启用GPU加速（若设备支持） #

将所有张量移至CUDA GPU上计算，GPU的并行计算能力会大幅提升大张量的运算速度，尤其当n较大时效果显著：

# 在张量初始化后添加该行
x = x.cuda()

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