实验记录 训练过程中网络参数的更新梯度爆炸，出现nan值

一、遇到的问题

训练的过程中，网络的输出突然为nan值， 导致训练终止。

二、进行排查

                    loss.backward()
                    # ==========
                    # for name, param in self.actor.net.named_parameters():
                    #     if param.grad is not None:
                    #         if 'state' in name:
                    #             print("prompt_embeddings_state 梯度统计:")
                    #             grad = param.grad
                    #             print(f"Mean: {grad.mean().item():.6f}")
                    #             print(f"Max: {grad.max().item():.6f}")
                    #             print(f"Min: {grad.min().item():.6f}")
                    #             print(f"Std: {grad.std().item():.6f}")
                    #         grad_has_nan = torch.isnan(param.grad).any()  # 检查梯度是否有 NaN
                    #         if grad_has_nan:
                    #             has_nan_grad = True
                    #             print(f"⚠️ 梯度 NaN 检测: 模块 '{name}' 的梯度包含 NaN！")
                    #             print(f"   梯度形状: {param.grad.shape}")
                    #             print(f"   NaN 数量: {torch.isnan(param.grad).sum().item()} 个")
                    #             print(f"   梯度示例:\n{param.grad}\n")  # 打印梯度详情
                    #         # print(f"Gradient of {name}:")
                    #         # print(param.grad)  # 打印梯度张量
                    #         # print(f"Gradient shape: {param.grad.shape}")
                    #     #     else:
                    #     #         print(f"✅ 模块 '{name}' 的梯度正常（无 NaN）")
                    #     # else:
                    #     #     print(f"ℹ️ 模块 '{name}' 无梯度（可能是冻结参数或未参与计算）")

加在loss.backward()后，去排查那块参数出现了问题。

打印的信息

⚠️ 梯度 NaN 检测: 模块 'backbone.prompt_embeddings_state' 的梯度包含 NaN！
   梯度形状: torch.Size([1, 40, 768])
   NaN 数量: 30720 个
   梯度示例:
tensor([[[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         ...,
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]]], device='cuda:0')

然后检查与其相关的前向计算。

在我的训练过程中，因为我要计算kl散度，当时的实现
 

# state_kl_loss = F.kl_div(
        #     input=F.log_softmax(s1, dim=-1),  # 输入分布（log_softmax）
        #     target=F.softmax(s2, dim=-1),  # 目标分布（softmax）
        #     reduction='batchmean'  # 对批次取平均
        # )

各种打印信息
 

 # 7. 计算 s1 和 s2 的 KL 散度（KL(s1 || s2)）
        # 注意：KL 散度是非对称的，这里计算 KL(s1_avg || s2_avg)

        # print("s1 统计信息 (before log_softmax):")
        # print(f"Mean: {s1.mean().item():.6f}, Std: {s1.std().item():.6f}")
        # print(f"Max: {s1.max().item():.6f}, Min: {s1.min().item():.6f}")
        #
        # print("s2 统计信息 (before softmax):")
        # print(f"Mean: {s2.mean().item():.6f}, Std: {s2.std().item():.6f}")
        # print(f"Max: {s2.max().item():.6f}, Min: {s2.min().item():.6f}")

        # log_softmax_s1 = F.log_softmax(s1, dim=-1)
        # softmax_s2 = F.softmax(s2, dim=-1)
        #
        # print("log_softmax(s1) 统计信息:")
        # print(f"Mean: {log_softmax_s1.mean().item():.6f}, Std: {log_softmax_s1.std().item():.6f}")
        # print(f"Max: {log_softmax_s1.max().item():.6f}, Min: {log_softmax_s1.min().item():.6f}")
        # print("是否有 -inf:", (log_softmax_s1 == -float('inf')).any())
        #
        # print("softmax(s2) 统计信息:")
        # print(f"Mean: {softmax_s2.mean().item():.6f}, Std: {softmax_s2.std().item():.6f}")
        # print(f"Max: {softmax_s2.max().item():.6f}, Min: {softmax_s2.min().item():.6f}")
        # print("是否有 0:", (softmax_s2 == 0.0).any())
        # print("是否有 inf:", torch.isinf(softmax_s2).any())

发现
 

s1 统计信息 (before log_softmax):
Mean: -0.162447, Std: 3.158298
Max: 17.990225, Min: -48.743835
s2 统计信息 (before softmax):
Mean: -0.157571, Std: 5.495389
Max: 40.220631, Min: -63.233921
prompt_embeddings_state 梯度统计:

Mean: -0.000205
Max: 8.615377
Min: -6.830388
Std: 0.324599

下一轮之后就变成了
Mean: nan
Max: nan
Min: nan
Std: nan

这是因为要把tensor变成概率分布，进行了取softmax操作，这之后，最小值就有可能非常接近0了，因为数据的维度可是3万多，这么多的数据一起进行softmax，加和为1，按平均来算每个还不过0.00003

具体的打印了一下，发现
 

印信息
log_softmax(s1) 统计信息:
Mean: -20.639292, Std: 3.170919
Max: -2.767323, Min: -69.401436
是否有 -inf: tensor(False, device='cuda:0')
softmax(s2) 统计信息:
Mean: 0.000033, Std: 0.001325
Max: 0.117576, Min: 0.000000
是否有 0: tensor(True, device='cuda:0')
是否有 inf: tensor(False, device='cuda:0')
prompt_embeddings_state 梯度统计:
Mean: nan
Max: nan
Min: nan
Std: nan

确实，softmax之后的最小值为0

三、关于kl散度的问题

用上述的方法实现kl散度确实存在这个问题，那我们可以换个角度实现。

state_kl_loss = -torch.sum(F.log_softmax(s1, dim=1) * F.softmax(s2, dim=1), dim=1).mean()

该函数实现

与kl散度的关系

默认 

自身的信息熵为常数 

四、yolov9训练时遇到的梯度爆炸问题
 

在对yolov9进行改进时，添加了设计的额外模块，导致刚开始训练，梯度就出现了nan值。在研究之后，发现是其使用的混合精度训练导致的该问题，因此，需要取消这种混合精度训练。