mode=args.eval_file_prefix, is_training=False, ) # Create callbacks for train and eval modes train_callback = nemo.core.SimpleLossLoggerCallback( tensors=train_tensors, print_func=lambda x: logging.info("Loss: {:.3f}".format(x[0].item())), tb_writer=nf.tb_writer, get_tb_values=lambda x: [["loss", x[0]]], step_freq=args.loss_step_freq, ) eval_callback = nemo.core.EvaluatorCallback( eval_tensors=eval_tensors, user_iter_callback=lambda x, y: eval_iter_callback(x, y, data_layer), user_epochs_done_callback=lambda x: eval_epochs_done_callback( x, f'{nf.work_dir}/graphs'), tb_writer=nf.tb_writer, eval_step=steps_per_epoch, ) # Create callback to save checkpoints 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)
val_loss = loss(logits=val_logits, labels=val_labels) train_data_size = len(train_data) steps_per_epoch = math.ceil(train_data_size / (BATCH_SIZE * NUM_GPUS)) train_callback = nemo.core.SimpleLossLoggerCallback( tensors=[train_loss, train_logits], print_func=lambda x: nemo.logging.info( f'Train loss: {str(np.round(x[0].item(), 3))}'), tb_writer=nf.tb_writer, get_tb_values=lambda x: [["train_loss", x[0]]], step_freq=steps_per_epoch) eval_callback = nemo.core.EvaluatorCallback( eval_tensors=[val_logits, val_labels], user_iter_callback=lambda x, y: eval_iter_callback(x, y, val_data), user_epochs_done_callback=lambda x: eval_epochs_done_callback( x, 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=1) lr_policy_fn = get_lr_policy('WarmupAnnealing', total_steps=NUM_EPOCHS * steps_per_epoch, warmup_ratio=0.1) nf.train( tensors_to_optimize=[train_loss], callbacks=[train_callback, eval_callback, ckpt_callback],