D_wass_valid_epoch_avg,
G_cost_epoch_avg
)
)
# Generate audio samples.
if epoch % epochs_per_sample == 0:
LOGGER.info("Generating samples...")
sample_out = netG(sample_noise_Var)
if cuda:
sample_out = sample_out.cpu()
sample_out = sample_out.data.numpy()
save_samples(sample_out, epoch, output_dir)
LOGGER.info('>>>>>>>Training finished !<<<<<<<')
# Save model
LOGGER.info("Saving models...")
netD_path = os.path.join(output_dir, "discriminator.pkl")
netG_path = os.path.join(output_dir, "generator.pkl")
torch.save(netD.state_dict(), netD_path, pickle_protocol=pickle.HIGHEST_PROTOCOL)
torch.save(netG.state_dict(), netG_path, pickle_protocol=pickle.HIGHEST_PROTOCOL)
# Plot loss curve.
LOGGER.info("Saving loss curve...")
plot_loss(D_costs_train, D_wasses_train, D_costs_valid, D_wasses_valid, G_costs, output_dir)
LOGGER.info("All finished!")
lmbda=args['lmbda'],
use_cuda=ngpus >= 1,
discriminator_updates=args['discriminator_updates'],
latent_dim=latent_dim,
epochs_per_sample=args['epochs_per_sample'],
sample_size=args['sample_size'])
print('Finished training.')
print('Final discriminator loss on validation and test:')
# LOGGER.info(pprint.pformat(final_discr_metrics))
print('Saving models...')
model_gen_output_path = os.path.join(model_dir, "model_gen.pkl")
model_dis_output_path = os.path.join(model_dir, "model_dis.pkl")
torch.save(model_gen.state_dict(),
model_gen_output_path,
pickle_protocol=pk.HIGHEST_PROTOCOL)
torch.save(model_dis.state_dict(),
model_dis_output_path,
pickle_protocol=pk.HIGHEST_PROTOCOL)
print('Saving metrics...')
history_output_path = os.path.join(model_dir, "history.pkl")
final_discr_metrics_output_path = os.path.join(model_dir,
"final_discr_metrics.pkl")
with open(history_output_path, 'wb') as f:
pk.dump(history, f)
with open(final_discr_metrics_output_path, 'wb') as f:
pk.dump(final_discr_metrics, f)
output_dir=model_dir,
lmbda=args['lmbda'],
use_cuda=ngpus>=1,
discriminator_updates=args['discriminator_updates'],
latent_dim=latent_dim,
epochs_per_sample=args['epochs_per_sample'],
sample_size=args['sample_size'])
LOGGER.info('Finished training.')
LOGGER.info('Final discriminator loss on validation and test:')
LOGGER.info(pprint.pformat(final_discr_metrics))
LOGGER.info('Saving models...')
model_gen_output_path = os.path.join(model_dir, "model_gen.pkl")
model_dis_output_path = os.path.join(model_dir, "model_dis.pkl")
torch.save(model_gen.state_dict(), model_gen_output_path,
pickle_protocol=pk.HIGHEST_PROTOCOL)
torch.save(model_dis.state_dict(), model_dis_output_path,
pickle_protocol=pk.HIGHEST_PROTOCOL)
LOGGER.info('Saving metrics...')
history_output_path = os.path.join(model_dir, "history.pkl")
final_discr_metrics_output_path = os.path.join(model_dir, "final_discr_metrics.pkl")
with open(history_output_path, 'wb') as f:
pk.dump(history, f)
with open(final_discr_metrics_output_path, 'wb') as f:
pk.dump(final_discr_metrics, f)
LOGGER.info('Done!')
