PyTorch中能否根据批量大小调整优化器步长？

当训练时最后一批样本数量小于标准批次大小（比如标准64，最后一批6个），常规的梯度更新会让小批次的权重调整幅度与大批次一致，但实际上每个样本的贡献应该相同，因此需要对小批次的梯度更新进行缩放。以下是三种实现方式：

方法1：动态缩放梯度 #

反向传播完成后，遍历模型所有参数，将梯度乘以当前批次大小与标准批次大小的比值，再执行优化器更新。这样能保证每个样本对权重更新的贡献一致。

示例代码：

standard_batch_size = 64
current_batch_size = len(target)  # 获取当前批次的样本数量

optimizer.zero_grad()
loss.backward()

# 对梯度进行缩放
for param in model.parameters():
    if param.grad is not None:
        param.grad.data.mul_(current_batch_size / standard_batch_size)

optimizer.step()

方法2：调整损失函数的计算逻辑 #

将损失函数的计算从默认的平均损失改为总损失除以标准批次大小，这样反向传播得到的梯度天然就是按标准批次缩放后的结果，无需额外处理梯度。

示例代码（以交叉熵损失为例）：

standard_batch_size = 64
# 将损失函数的reduction设置为sum，计算总损失后除以标准批次大小
criterion = nn.CrossEntropyLoss(reduction='sum')

optimizer.zero_grad()
output = model(input)
loss = criterion(output, target) / standard_batch_size
loss.backward()
optimizer.step()

方法3：动态调整优化器学习率 #

临时将优化器的学习率乘以当前批次大小与标准批次大小的比值，执行更新后再恢复原学习率。这种方法适合不想修改损失计算或梯度的场景。

示例代码：

standard_batch_size = 64
current_batch_size = len(target)
lr_scale = current_batch_size / standard_batch_size

# 保存原学习率
original_lrs = [pg['lr'] for pg in optimizer.param_groups]

# 调整学习率
for pg in optimizer.param_groups:
    pg['lr'] *= lr_scale

optimizer.zero_grad()
loss.backward()
optimizer.step()

# 恢复原学习率
for i, pg in enumerate(optimizer.param_groups):
    pg['lr'] = original_lrs[i]

注意事项 #

• 若损失函数原本就是按总损失计算（reduction='sum'），只需确保除以标准批次大小而非当前批次大小即可。
• 方法1和方法2本质逻辑一致，都是将每个样本的梯度贡献标准化到标准批次水平，可根据代码习惯选择。

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