}, os.path.join(args.save, 'checkpt.pth'))
model.train()
if itr % args.viz_freq == 0:
with torch.no_grad():
model.eval()
p_samples = toy_data.inf_train_gen(args.data, batch_size=2000)
sample_fn, density_fn = get_transforms(model)
plt.figure(figsize=(9, 3))
visualize_transform(
p_samples, torch.randn, standard_normal_logprob, transform=sample_fn, inverse_transform=density_fn,
samples=True, npts=800, device=device
)
fig_filename = os.path.join(args.save, 'figs', '{:04d}.pdf'.format(itr))
utils.makedirs(os.path.dirname(fig_filename))
plt.savefig(fig_filename)
plt.close()
model.train()
end = time.time()
logger.info('Training has finished.')
save_traj_dir = os.path.join(args.save, 'trajectory')
logger.info('Plotting trajectory to {}'.format(save_traj_dir))
data_samples = toy_data.inf_train_gen(args.data, batch_size=2000)
save_trajectory(model, data_samples, save_traj_dir, device=device)
trajectory_to_video(save_traj_dir)
standard_normal_logprob,
transform=sample_fn,
inverse_transform=density_fn,
samples=True,
npts=800,
device=device)
fig_filename = os.path.join(
args.save, 'figs',
'{:04d}_{:01d}.jpg'.format(itr, i + 1))
utils.makedirs(os.path.dirname(fig_filename))
plt.savefig(fig_filename)
plt.close()
model.train()
end = time.time()
logger.info('Training has finished.')
logger.info('Plotting trajectory')
save_traj_dir = os.path.join(args.save, 'trajectory')
logger.info('Plotting trajectory to {}'.format(save_traj_dir))
data_samples = sample_data(img2, batch_size=2000)
save_trajectory(model,
data_samples,
save_traj_dir,
device=device,
ntimes=100,
memory=0.01,
end_times=times)
trajectory_to_video(save_traj_dir)
