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)
nf.train(
tensors_to_optimize=[train_loss],
train_loss = loss(logits=train_logits, labels=train_labels)
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)
#ask if training is to be done
TRAIN = bool(input("Do you want to train? [True|False]"))
if TRAIN:
nf.train(tensors_to_optimize=[train_loss],
