如何在使用二元交叉熵损失的神经网络中计算Armijo步长？

在你的场景里，f(x)就是当前模型参数下的二元交叉熵损失值，f(x + lr*v)则是参数临时更新为x + lr*v后的损失值。以下是具体实现步骤和代码：

核心逻辑 #

• 先计算当前参数对应的损失f(x)
• 临时更新模型参数为x + lr*v，计算此时的损失f(x + lr*v)
• 计算完成后立刻恢复原参数，避免影响后续训练流程

具体代码示例 #

假设你已准备好：

• 训练用的输入inputs和标签targets
• 训练中的模型model
• 损失函数criterion = nn.BCELoss()
• 下降方向v（与模型参数结构完全匹配的参数组）
• 候选步长lr，Armijo系数c（通常取0.01~0.1）

import torch
import torch.nn as nn

# 1. 计算当前损失f(x)
model.eval()
with torch.no_grad():
    outputs = model(inputs)
    f_x = criterion(outputs, targets).item()

# 2. 保存原参数，用于后续恢复
original_params = [param.clone() for param in model.parameters()]

# 3. 临时更新参数为x + lr*v，计算f(x + lr*v)
with torch.no_grad():
    for param, delta in zip(model.parameters(), v):
        param.data.add_(lr * delta)  # 执行参数更新：param = param + lr*v

with torch.no_grad():
    outputs_new = model(inputs)
    f_x_lrv = criterion(outputs_new, targets).item()

# 4. 恢复原模型参数
with torch.no_grad():
    for param, orig_param in zip(model.parameters(), original_params):
        param.data.copy_(orig_param)

# 5. 验证Armijo条件（注意：原公式中的func_gradient应为梯度与下降方向v的点积<∇f(x), v>）
# 假设你已计算出当前参数的梯度grad（与模型参数结构匹配）
grad_dot_v = sum(torch.sum(g * d) for g, d in zip(grad, v)).item()
armijo_condition = f_x_lrv <= f_x + c * lr * grad_dot_v

关键细节提醒 #

• 所有参数操作和损失计算都要包裹在torch.no_grad()里，避免不必要的梯度计算，节省显存同时防止干扰后续训练的梯度流
• 下降方向v、梯度grad的结构必须和模型参数完全对齐（每个参数张量的形状、数量一致）
• 原公式里的func_gradient对应Armijo准则的标准形式应为梯度与下降方向的点积，需确认你使用的公式是否匹配该计算逻辑

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