cur_best_valid_recons = np.inf
# Set early stop
early_stop = 0
# Through the epochs
for epoch in range(1, args.epochs + 1, 1):
print(f"Epoch: {epoch}")
# Training epoch
loss_mean, kl_div_mean, recon_loss_mean = learn.train(
model, optimizer, criterion, args, epoch)
# Validate epoch
loss_mean_validate, kl_div_mean_validate, recon_loss_mean_validate = learn.validate(
model, criterion, args, epoch)
# Step for learning rate
scheduler.step(loss_mean_validate)
# Test model
loss_mean_test, kl_div_mean_test, recon_loss_mean_test = learn.test(
model, criterion, args, epoch)
# Compare input data and reconstruction
if (epoch % 25 == 0):
reconstruction(args, model, epoch, test_set)
# Gather losses
loss_list = [loss_mean, loss_mean_validate, loss_mean_test]
for counter, loss in enumerate(loss_list):
losses[epoch - 1, counter] = loss
# Gather reconstruction losses
recon_loss_list = [
recon_loss_mean, recon_loss_mean_validate, recon_loss_mean_test
]
for counter, loss in enumerate(recon_loss_list):
recon_losses[epoch - 1, counter] = loss
# Save losses
torch.save({
