TensorFlow转PyTorch：指定学习率衰减与Adam优化器代码转换

对应代码转换 #

原TensorFlow代码实现了指数阶梯式衰减学习率搭配Adam优化器，以下是等价的PyTorch实现：

首先定义基础优化器，再配置学习率调度器：

import torch
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR

# 假设start_lr是预先定义的初始学习率，比如start_lr=0.001
optimizer = Adam(model.parameters(), lr=start_lr, eps=0.1)
# StepLR对应TensorFlow中staircase=True的exponential_decay
# step_size=3000 对应原代码的decay_steps，gamma=0.96对应decay_rate
scheduler = StepLR(optimizer, step_size=3000, gamma=0.96)

关键细节说明 #

• TensorFlow的tf.train.exponential_decay设置staircase=True时，是每3000步将学习率乘以0.96，这和PyTorch的StepLR行为完全一致。
• Adam优化器的参数对应：TensorFlow的epsilon等价于PyTorch的eps。
• 使用时，在每个训练步骤更新参数后，需要调用scheduler.step()来更新学习率：

# 训练循环示例
for step in range(total_steps):
    # 前向传播、计算损失
    loss = model(inputs)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    # 更新学习率
    scheduler.step()

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