def train(param, x, y):
dim_in = x.shape[1]
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
dataloader = DataLoader(torch.from_numpy(x.astype(np.float32)),
batch_size=param.batch_size,
shuffle=True,
num_workers=2)
flow = RealNVP(dim_in, device)
flow.to(device)
flow.train()
optimizer = torch.optim.Adam(
[p for p in flow.parameters() if p.requires_grad == True], lr=param.lr)
it, print_cnt = 0, 0
while it < param.total_it:
for i, data in enumerate(dataloader):
loss = -flow.log_prob(data.to(device)).mean()
optimizer.zero_grad()
loss.backward(retain_graph=True)
optimizer.step()
it += data.shape[0]
print_cnt += data.shape[0]
if print_cnt > PRINT_FREQ:
print('it {:d} -- loss {:.03f}'.format(it, loss))
print_cnt = 0
torch.save(flow.state_dict(), 'flow_model.pytorch')
bst = xgb.train(param, xgb.DMatrix(x[:num_train], label=y[:num_train]),
num_round)
pred_val = bst.predict(xgb.DMatrix(x[num_train:]))
val_auc = roc_auc_score(y[num_train:], pred_val)
print('\nAUC score on val set: {:.03f}'.format(val_auc))
pred_fn = lambda x: bst.predict(xgb.DMatrix(x))
shap_fn = lambda x: bst.predict(xgb.DMatrix(x), pred_contribs=True)
inf_fn = lambda x: flow.f(
torch.from_numpy(x.astype(np.float32)).to(device))[0].detach().cpu(
).numpy()
gen_fn = lambda z: flow.g(
torch.from_numpy(z.astype(np.float32)).to(device)).detach().cpu(
).numpy()
logp_fn = lambda x: flow.log_prob(
torch.from_numpy(x.astype(np.float32)).to(device)).detach().cpu(
).numpy()
if args.sensitivity:
noise_sd = 0.1
n_nbhrs = 40
i = np.random.randint(num_train, x.shape[0], 1)[0]
print('\nSensitivity for sample {:d}'.format(i))
x_test = x[i][None, :]
z_test = inf_fn(x_test)
z_nbhr = z_test + noise_sd * np.random.randn(
n_nbhrs, z_test.shape[1]).astype(np.float32)
x_nbhr = gen_fn(z_nbhr)