valid=0.0,
batch_size=batch_size,
num_workers=0)
best_loss, best_label, best_pred, best_mask, best_pear = None, None, None, None, None
# Training loop
for e in tqdm(range(n_epochs), desc='MOSEI Regression'):
train_loss, train_mae, train_pear, _, _, _, _ = train_or_eval_model(
model, loss_function, train_loader, e, optimizer, True, cuda=cuda)
test_loss, test_mae, test_pear, test_label, test_pred, test_mask, _ = train_or_eval_model(
model, loss_function, test_loader, e, cuda=cuda)
writer.add_scalar("Train Loss - MOSEI Regression", train_loss, e)
writer.add_scalar("Test Loss - MOSEI Regression", test_loss, e)
writer.add_scalar("Train MAE - MOSEI Regression", train_mae, e)
writer.add_scalar("Test MAE - MOSEI Regression", test_mae, e)
writer.add_scalar("Train Pearson - MOSEI Regression", train_pear, e)
writer.add_scalar("Test Pearson - MOSEI Regression", test_pear, e)
if best_loss == None or best_loss > test_loss:
best_loss, best_label, best_pred, best_mask, best_pear =\
test_loss, test_label, test_pred, test_mask, test_pear
torch.save(model.state_dict(), args.model_path)
print('Model saved at {}'.format(args.model_path), file=sys.stderr)
print('Test performance..')
print('Loss {} MAE {} r {}'.format(
best_loss,
round(
mean_absolute_error(best_label, best_pred,
sample_weight=best_mask), 4), best_pear))
