def main():
if args.data == 'boston':
data = np.load('data_small/boston_no_discrete.npy')
elif args.data == 'white_wine':
data = np.load('data_small/white_no_discrete_no_corr_0.98.npy')
elif args.data == 'red_wine':
data = np.load('data_small/red_no_discrete_no_corr_0.98.npy')
n_features = data.shape[1]
n_train = int(data.shape[0] * 0.9)
train_data_clean = util_shuffle(data[:n_train])
test_data = data[:n_train]
kf = KFold(n_splits=10)
covar = np.diag(args.covar * np.ones((n_features, )))
train_data = train_data_clean + \
np.random.multivariate_normal(mean=np.zeros(
(n_features,)), cov=covar, size=n_train)
# train_covars = np.repeat(
# covar[np.newaxis, :, :], n_train, axis=0)
# train_dataset = DeconvDataset(train_data, train_covars)
for i, (train_index, eval_index) in enumerate(kf.split(train_data)):
X_train, X_eval = train_data[train_index], train_data[eval_index]
train_covars = np.repeat(covar[np.newaxis, :, :],
X_train.shape[0],
axis=0)
eval_covars = np.repeat(covar[np.newaxis, :, :],
X_eval.shape[0],
axis=0)
train_dataset = DeconvDataset(X_train, train_covars)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
eval_dataset = DeconvDataset(X_eval, eval_covars)
eval_loader = DataLoader(eval_dataset,
batch_size=args.batch_size,
shuffle=False)
model = SVIFlowToy(dimensions=n_features,
objective=args.objective,
posterior_context_size=n_features,
batch_size=args.batch_size,
device=device,
maf_steps_prior=args.flow_steps_prior,
maf_steps_posterior=args.flow_steps_posterior,
maf_features=args.maf_features,
maf_hidden_blocks=args.maf_hidden_blocks,
K=args.K)
message = 'Total number of parameters: %s' % (sum(
p.numel() for p in model.parameters()))
logger.info(message)
optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr)
# training
scheduler = [30]
epoch = 0
best_model = copy.deepcopy(model.state_dict())
best_eval_loss = compute_eval_loss(model, eval_loader, device,
len(eval_index))
n_epochs_not_improved = 0
model.train()
while n_epochs_not_improved < scheduler[-1] and epoch < args.n_epochs:
for batch_idx, data in enumerate(train_loader):
data[0] = data[0].to(device)
data[1] = data[1].to(device)
for prior_params in model.model._prior.parameters():
prior_params.requires_grad = True
for post_params in model.model._approximate_posterior.parameters(
):
post_params.requires_grad = False
for i in range(args.prior_iter):
loss = -model.score(data).mean()
message = 'Loss prior %s: %f' % (i, loss)
logger.info(message)
optimizer.zero_grad()
loss.backward(retain_graph=True)
optimizer.step()
for prior_params in model.model._prior.parameters():
prior_params.requires_grad = False
for post_params in model.model._approximate_posterior.parameters(
):
post_params.requires_grad = True
for i in range(args.posterior_iter):
loss = -model.score(data).mean()
message = 'Loss posterior %s: %f' % (i, loss)
logger.info(message)
optimizer.zero_grad()
loss.backward(retain_graph=True)
optimizer.step()
model.eval()
test_loss_clean = - \
model.model._prior.log_prob(
torch.from_numpy(test_data).to(device)).mean()
message = 'Test loss (clean) = %.5f' % test_loss_clean
logger.info(message)
eval_loss = compute_eval_loss(model, eval_loader, device,
len(eval_index))
if eval_loss < best_eval_loss:
best_model = copy.deepcopy(model.state_dict())
best_eval_loss = eval_loss
n_epochs_not_improved = 0
else:
n_epochs_not_improved += 1
model.train()
epoch += 1
break
model = model.load_state_dict(best_model)
test_loss_clean = - \
model.model._prior.log_prob(
torch.from_numpy(test_data).to(device)).mean()
message = 'Final test loss (clean) = %.5f' % test_loss_clean
logger.info(message)
def main():
train_covar = covar_gen(args.covar, args.n_train_points).astype(np.float32)
train_data_clean = data_gen(args.data,
args.n_train_points)[0].astype(np.float32)
# plt.scatter(train_data_clean[:, 0], train_data_clean[:, 1])
train_data = np.zeros_like(train_data_clean)
for i in range(args.n_train_points):
train_data[i] = train_data_clean[i] + np.random.multivariate_normal(
mean=np.zeros((2, )), cov=train_covar[i])
# plt.scatter(train_data[:, 0], train_data[:, 1])
# plt.show()
train_covar = torch.from_numpy(train_covar)
train_data = torch.from_numpy(train_data.astype(np.float32))
train_dataset = DeconvDataset(train_data, train_covar)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
test_data_clean = torch.from_numpy(
data_gen(args.data, args.n_test_points)[0].astype(np.float32))
eval_covar = covar_gen(args.covar, args.n_eval_points).astype(np.float32)
eval_data_clean = data_gen(args.data,
args.n_eval_points)[0].astype(np.float32)
eval_data = np.zeros_like(eval_data_clean)
for i in range(args.n_eval_points):
eval_data[i] = eval_data_clean[i] + np.random.multivariate_normal(
mean=np.zeros((2, )), cov=eval_covar[i])
eval_covar = torch.from_numpy(eval_covar)
eval_data = torch.from_numpy(eval_data.astype(np.float32))
eval_dataset = DeconvDataset(eval_data, eval_covar)
eval_loader = DataLoader(eval_dataset,
batch_size=args.test_batch_size,
shuffle=False)
if args.infer == 'true_data':
model = SVIFlowToy(dimensions=2,
objective=args.objective,
posterior_context_size=args.posterior_context_size,
batch_size=args.batch_size,
device=device,
maf_steps_prior=args.flow_steps_prior,
maf_steps_posterior=args.flow_steps_posterior,
maf_features=args.maf_features,
maf_hidden_blocks=args.maf_hidden_blocks,
K=args.K)
else:
model = SVIFlowToyNoise(
dimensions=2,
objective=args.objective,
posterior_context_size=args.posterior_context_size,
batch_size=args.batch_size,
device=device,
maf_steps_prior=args.flow_steps_prior,
maf_steps_posterior=args.flow_steps_posterior,
maf_features=args.maf_features,
maf_hidden_blocks=args.maf_hidden_blocks,
K=args.K)
message = 'Total number of parameters: %s' % (sum(
p.numel() for p in model.parameters()))
logger.info(message)
optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr)
#training
scheduler = list(map(int, args.eval_based_scheduler.split(',')))
epoch = 0
best_model = copy.deepcopy(model.state_dict())
best_eval_loss = compute_eval_loss(model, eval_loader, device,
args.n_eval_points)
n_epochs_not_improved = 0
model.train()
while n_epochs_not_improved < scheduler[-1] and epoch < args.n_epochs:
for batch_idx, data in enumerate(train_loader):
data[0] = data[0].to(device)
data[1] = data[1].to(device)
loss = -model.score(data).mean()
optimizer.zero_grad()
loss.backward(retain_graph=True)
optimizer.step()
model.eval()
eval_loss = compute_eval_loss(model, eval_loader, device,
args.n_eval_points)
if eval_loss < best_eval_loss:
best_model = copy.deepcopy(model.state_dict())
best_eval_loss = eval_loss
n_epochs_not_improved = 0
else:
n_epochs_not_improved += 1
lr_scheduler(n_epochs_not_improved, optimizer, scheduler, logger)
if (epoch + 1) % args.test_freq == 0:
if args.infer == 'true_data':
test_loss_clean = -model.model._prior.log_prob(
test_data_clean.to(device)).mean()
else:
test_loss_clean = -model.model._likelihood.log_prob(
test_data_clean.to(device)).mean()
message = 'Epoch %s:' % (
epoch + 1
), 'train loss = %.5f' % loss, 'eval loss = %.5f' % eval_loss, 'test loss (clean) = %.5f' % test_loss_clean
logger.info(message)
else:
message = 'Epoch %s:' % (
epoch +
1), 'train loss = %.5f' % loss, 'eval loss = %.5f' % eval_loss
logger.info(message)
if (epoch + 1) % args.viz_freq == 0:
if args.infer == 'true_data':
samples = model.model._prior.sample(
1000).detach().cpu().numpy()
else:
samples = model.model._likelihood.sample(
1000).detach().cpu().numpy()
corner.hist2d(samples[:, 0], samples[:, 1])
fig_filename = args.dir + 'out/' + name + '_corner_fig_' + str(
epoch + 1) + '.png'
plt.savefig(fig_filename)
plt.close()
plt.scatter(samples[:, 0], samples[:, 1])
fig_filename = args.dir + 'out/' + name + '_scatter_fig_' + str(
epoch + 1) + '.png'
plt.savefig(fig_filename)
plt.close()
model.train()
epoch += 1
model.load_state_dict(best_model)
model.eval()
if args.infer == 'true_data':
test_loss_clean = -model.model._prior.log_prob(
test_data_clean.to(device)).mean()
else:
test_loss_clean = -model.model._likelihood.log_prob(
test_data_clean.to(device)).mean()
message = 'Final test loss (clean) = %.5f' % test_loss_clean
logger.info(message)
torch.save(model.state_dict(), args.dir + 'models/' + name + '.model')
logger.info('Training has finished.')
if args.data.split('_')[0] == 'mixture' or args.data.split(
'_')[0] == 'gaussian':
kl_points = data_gen(args.data, args.n_kl_points)[0].astype(np.float32)
if args.infer == 'true_data':
model_log_prob = model.model._prior.log_prob(
torch.from_numpy(kl_points.astype(
np.float32)).to(device)).mean()
else:
model_log_prob = model.model._likelihood.log_prob(
torch.from_numpy(kl_points.astype(
np.float32)).to(device)).mean()
data_log_prob = compute_data_ll(args.data, kl_points).mean()
approximate_KL = data_log_prob - model_log_prob
message = 'KL div %.5f:' % approximate_KL
logger.info(message)