def get_ckpt_model_and_data(args):
# Load checkpoint.
checkpt = torch.load(args.checkpt,
map_location=lambda storage, loc: storage)
ckpt_args = checkpt['args']
state_dict = checkpt['state_dict']
# Construct model and restore checkpoint.
regularization_fns, regularization_coeffs = create_regularization_fns(
ckpt_args)
model = build_model_tabular(ckpt_args, 2,
regularization_fns).to(device)
if ckpt_args.spectral_norm: add_spectral_norm(model)
set_cnf_options(ckpt_args, model)
model.load_state_dict(state_dict)
model.to(device)
print(model)
print("Number of trainable parameters: {}".format(
count_parameters(model)))
# Load samples from dataset
data_samples = toy_data.inf_train_gen(ckpt_args.data, batch_size=2000)
return model, data_samples
def compute_loss(args, model, batch_size=None):
if batch_size is None: batch_size = args.batch_size
# load data
x = toy_data.inf_train_gen(args.data, batch_size=batch_size)
x = torch.from_numpy(x).type(torch.float32).to(device)
zero = torch.zeros(x.shape[0], 1).to(x)
# transform to z
z, delta_logp = model(x, zero)
# compute log q(z)
logpz = standard_normal_logprob(z).sum(1, keepdim=True)
logpx = logpz - delta_logp
loss = -torch.mean(logpx)
return loss
log_message = '[TEST] Iter {:04d} | Test Loss {:.6f}'.format(itr, test_loss)
logger.info(log_message)
if test_loss.item() < best_loss:
best_loss = test_loss.item()
utils.makedirs(args.save)
torch.save({
'args': args,
'state_dict': model.state_dict(),
}, 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}.jpg'.format(itr))
utils.makedirs(os.path.dirname(fig_filename))
plt.savefig(fig_filename)
plt.close()
model.train()
end = time.time()