INNER CODE UNIT · Python
lr
yixinL7/BRIO · main.py:453
lr = args.max_lr * min(all_step_cnt ** (-0.5), all_step_cnt * (args.warmup_steps ** (-1.5)))
for param_group in s_optimizer.param_groups:
param_group['lr'] = lr
s_optimizer.step()
s_optimizer.zero_grad()
if epoch_step % args.report_freq == 0 and step_cnt == 0 and is_master:
# report stats
print("id: %d"%id)
print(f"similarity: {similarity[:, :10]}")
if not args.no_gold:
print(f"gold similarity: {gold_similarity}")
recorder.print("epoch: %d, batch: %d, avg loss: %.6f, avg ranking loss: %.6f, avg mle loss: %.6f"
%(epoch+1, epoch_step, avg_loss / args.report_freq, avg_ranking_loss / args.report_freq, avg_mle_loss / args.report_freq))
recorder.print(f"learning rate: {lr:.6f}")
recorder.plot("loss", {"loss": avg_loss / args.report_freq}, all_step_cnt)
recorder.plot("mle_loss", {"loss": avg_mle_loss / args.report_freq}, all_step_cnt)
recorder.plot("ranking_loss", {"loss": avg_ranking_loss / args.report_freq}, all_step_cnt)
recorder.print()