INNER CODE UNIT · Python
loss
yixinL7/BRIO · main.py:439
loss = args.rank_weight * ranking_loss + args.mle_weight * mle_loss
loss = loss / args.accumulate_step
avg_loss += loss.item()
avg_mle_loss += mle_loss.item() / args.accumulate_step
avg_ranking_loss += ranking_loss.item() / args.accumulate_step
loss.backward()
if step_cnt == args.accumulate_step:
# updating
if args.grad_norm > 0:
nn.utils.clip_grad_norm_(model.parameters(), args.grad_norm)
step_cnt = 0
epoch_step += 1
all_step_cnt += 1
# adjust learning rate
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()