val_loss = val_loss / len(val_dataloader) all_val_loss.append(val_loss) # Print Metrics print( f"\nEpoch: {epoch+1}/{epochs}:\ \nTrain Loss = {train_loss}\ \nVal Loss = {val_loss}" ) ################################################################################ ################################## Save & Test ################################# ################################################################################ # Generate Loss Graph graphics.draw_loss(all_train_loss, all_val_loss, output_prefix, mode="autoencoder") # Save Model torch.save( { "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "train_loss": all_train_loss, "val_loss": all_val_loss, "label_encoder": label_encoder, }, model_output_path, ) # Evaluation Time
train_loss = train_loss / len(train_dataloader) all_train_loss.append(train_loss) val_loss = val_loss / len(val_dataloader) all_val_loss.append(val_loss) # Print Metrics print(f"\nEpoch: {epoch+1}/{epochs}:\ \nTrain Loss = {train_loss}\ \nVal Loss = {val_loss}") ################################################################################ ################################## Save & Test ################################# ################################################################################ # Generate Loss Graph graphics.draw_loss(all_train_loss, all_val_loss, loss_output_path, mode="ae") # Save Model torch.save(model.state_dict(), model_output_path) # Plot Animation Sample fig, axis = graphics.make_grid(("Sample", sample), 4, 4) plt.savefig(animation_sample_image_name) # Create & Save Animation Writer = animation.writers["ffmpeg"] writer = Writer() anim = graphics.make_animation(graphics.make_grid, all_samples) anim.save(animation_output_path, writer=writer) model.eval()
netD.eval() # Generate a batch of fake images with torch.no_grad(): # Generate batch of latent vectors noise = torch.randn(batch_size, latent_dim, 1, 1, device=device) # Generate fake image batch with G generated = netG(noise).detach().cpu() # Plot fake images fig, axis = graphics.make_grid(("Test Sample", generated), 4, 4) plt.savefig(os.path.join(output_dir, f"epoch_{epoch}_test_sample_output.jpg")) ################################################################################ ################################## Save & Test ################################# ################################################################################ # Save output graphs graphics.draw_loss( all_generator_loss, all_discriminator_loss, loss_output_path, mode="gan" ) # Create & Save Animation anim = graphics.make_gan_animation(training_samples) anim.save(animation_output_path) # Pick a couple of sample images for an Input v Output comparison sample = data.get_samples_from_data(dataset, 16) # Plot A Set of Test Images fig, axis = graphics.make_grid(("Test Sample", sample), 4, 4) plt.savefig(test_sample_input_name)