print()
print("Fold:", fold)
print()
fold_checkpoints = os.path.join(experiment.checkpoints,
"fold{}".format(fold))
model, optimizer = get_model_optimizer(original_args)
checkpoint = os.path.join(fold_checkpoints, args.checkpoint)
state_dict = torch.load(checkpoint)
model.load_state_dict(state_dict)
del state_dict
torch.cuda.empty_cache()
test_preds = infer(original_args,
model,
test_loader,
test_shape=len(test_set))
test_preds_df = test_df[["qa_id"]].copy()
for k, col in enumerate(target_columns):
test_preds_df[col] = test_preds[:, k].astype(np.float32)
test_preds_df.to_csv(
os.path.join(args.output_dir, "fold-{}.csv".format(fold)),
index=False,
)
torch.cuda.empty_cache()
writer.writerow([epoch + 1, avg_loss, avg_val_loss, score])
torch.save(
model.state_dict(),
os.path.join(fold_checkpoints,
"model_on_epoch_{}.pth".format(epoch)),
)
val_preds_df = val_fold_df.copy()[["qa_id"] + args.target_columns]
val_preds_df[args.target_columns] = val_preds
val_preds_df.to_csv(
os.path.join(fold_predictions,
"val_on_epoch_{}.csv".format(epoch)),
index=False,
)
test_preds = infer(args, model, test_loader, test_shape=len(test_set))
test_preds_df = submission.copy()
test_preds_df[args.target_columns] = test_preds
test_preds_df.to_csv(
os.path.join(fold_predictions,
"test_on_epoch_{}.csv".format(epoch)),
index=False,
)
if score > best_score:
best_score = score
torch.save(
model.state_dict(),
os.path.join(fold_checkpoints, "best_model.pth"),
)
val_preds_df.to_csv(os.path.join(fold_predictions, "best_val.csv"),
