if args.wandb_project is not None:
wand_callback = nemo.core.WandbCallback(
train_tensors=[losses[0]],
wandb_name=args.wandb_exp_name,
wandb_project=args.wandb_project,
update_freq=args.loss_log_freq
if args.loss_log_freq > 0 else steps_per_epoch,
args=args,
)
callbacks.append(wand_callback)
eval_callback = nemo.core.EvaluatorCallback(
eval_tensors=eval_tensors,
user_iter_callback=lambda x, y: eval_iter_callback(x, y),
user_epochs_done_callback=lambda x: eval_epochs_done_callback(
x, punct_label_ids, capit_label_ids, graph_dir),
tb_writer=nf.tb_writer,
eval_step=args.eval_epoch_freq * steps_per_epoch,
wandb_name=args.wandb_exp_name,
wandb_project=args.wandb_project,
)
callbacks.append(eval_callback)
lr_policy_fn = get_lr_policy(args.lr_policy,
total_steps=args.num_epochs * steps_per_epoch,
warmup_ratio=args.lr_warmup_proportion)
nf.train(
tensors_to_optimize=[losses[0]],
callbacks=callbacks,
lr_policy=lr_policy_fn,
)
logging.info(f"steps_per_epoch = {steps_per_epoch}")
# Create trainer and execute training action
train_callback = nemo.core.SimpleLossLoggerCallback(
tensors=losses + train_logits,
print_func=lambda x: logging.info("Loss: {:.3f}".format(x[0].item())),
get_tb_values=lambda x: [["loss", x[0]]],
tb_writer=nf.tb_writer,
)
eval_callback = nemo.core.EvaluatorCallback(
eval_tensors=eval_tensors,
user_iter_callback=lambda x, y: eval_iter_callback(x, y),
user_epochs_done_callback=lambda x: eval_epochs_done_callback(
x, punct_label_ids, capit_label_ids, f'{nf.work_dir}/graphs'),
tb_writer=nf.tb_writer,
eval_step=steps_per_epoch,
)
ckpt_callback = nemo.core.CheckpointCallback(folder=nf.checkpoint_dir,
epoch_freq=args.save_epoch_freq,
step_freq=args.save_step_freq)
lr_policy_fn = get_lr_policy(args.lr_policy,
total_steps=args.num_epochs * steps_per_epoch,
warmup_ratio=args.lr_warmup_proportion)
nf.train(
tensors_to_optimize=[losses[0]],
callbacks=[train_callback, eval_callback, ckpt_callback],