train_X, train_Y, batch_size=batch_size)
val_batches_X, val_batches_Y, \
val_batches_counter_X, val_batches_counter_Y = bucket_and_batch(
val_X, val_Y, batch_size=batch_size)
test_batches_X, test_batches_Y,\
test_batches_counter_X, test_batches_counter_Y = bucket_and_batch(
test_X, test_Y, batch_size=batch_size)
print("hello")
test_accs = []
val_accs = []
for epoch in range(flags.epochs):
print("Training")
pretty_print('epoch', 'iter', 'loss', 'acc', 'penalty')
train_loss, train_acc, train_penalty = \
run(train_batches_X, train_batches_Y,
train_batches_counter_X, train_batches_counter_Y, train=True, epoch=epoch)
print("\n\n")
print("Validating")
pretty_print('epoch', 'iter', 'loss', 'acc', 'penalty')
val_loss, val_acc, val_penalty = \
run(val_batches_X, val_batches_Y,
val_batches_counter_X, val_batches_counter_Y, train=False, epoch=epoch)
print("\n\n")
print("Testing")
def run(batches_X, batches_Y, batches_counter_X, batches_counter_Y, train=True, epoch=0):
global model
global display_step
global global_step
if train:
model_ = model.train()
else:
model_ = model.eval()
total_loss = 0.0
total_penalty = 0.0
total_acc = 0.0
iter = 0
for batch_X, batch_Y, \
batch_counter_X, batch_counter_Y in zip(batches_X, batches_Y,
batches_counter_X, batches_counter_Y):
z, logits, batch_recon_X, batch_mu, batch_logvar = model_(batch_X)
labels = batch_Y
penalty = penalty_fn(logits, labels, train=train)
ER_penalty = ER_penalty_fn(z, labels)
acc = mean_accuracy(logits, labels)
nll = nll_fn(logits, labels)
if flags.ELBO_weight > 0:
save_image = True
else:
save_image = False
elbo = ELBO(batch_recon_X, batch_X, batch_mu, batch_logvar,
epoch=epoch, iter=iter, save_image_decision=save_image, model='vanilla')
weight_norm = T.tensor(0.).cuda()
for w in model.parameters():
weight_norm += w.norm().pow(2)
# print(elbo.item())
loss = nll*(1-flags.ELBO_weight) + elbo*flags.ELBO_weight
loss = loss + flags.ER_penalty_weight*ER_penalty
loss = loss + flags.l2 * weight_norm
penalty_weight = (flags.penalty_weight
if global_step >= flags.penalty_anneal_iters else 1.0)
loss = loss + penalty_weight * penalty
if penalty_weight > 1.0:
loss = loss/penalty_weight
pure_loss = nll.item()*(1-flags.ELBO_weight) + elbo.item()*flags.ELBO_weight
if train:
optimizer.zero_grad()
loss.backward()
optimizer.step()
if iter % display_step == 0:
pretty_print(
np.int32(epoch),
np.int32(iter),
np.float32(pure_loss),
acc.detach().cpu().numpy(),
penalty.detach().cpu().numpy(),
)
total_loss += pure_loss
total_acc += acc.item()
total_penalty += penalty.item()
if train:
global_step += 1
iter += 1
total_loss /= iter
total_acc /= iter
total_penalty /= iter
return total_loss, total_acc, total_penalty
def run(batches_X,
batches_Y,
batches_counter_X,
batches_counter_Y,
train=True,
epoch=0):
global model
global display_step
global global_step
global global_imap
if train:
global_imap = T.tensor(0.0).cuda()
model_ = model.train()
else:
model_ = model.eval()
total_loss = 0.0
total_penalty = 0.0
total_acc = 0.0
iter = 0
for batch_X, batch_Y, \
batch_counter_X, batch_counter_Y in zip(batches_X, batches_Y,
batches_counter_X, batches_counter_Y):
if train:
imap, imap_var, logits, noisy_logits, batch_recon_X, batch_mu, batch_logvar = model_(
batch_X, batch_counter_X)
else:
imap, imap_var, logits, noisy_logits, batch_recon_X, batch_mu, batch_logvar = model_(
batch_X, batch_counter_X, imap=global_imap)
labels = batch_Y
penalty = penalty_fn(logits, labels, train=train)
acc = mean_accuracy(logits, labels)
nll = nll_fn(logits, labels)
noisy_nll = nll_fn(noisy_logits, labels)
if flags.ELBO_weight > 0:
save_image = True
else:
save_image = False
elbo = ELBO(batch_recon_X,
batch_X,
batch_mu,
batch_logvar,
epoch=epoch,
iter=iter,
save_image_decision=save_image,
model='Causal_BVAE')
weight_norm_set_1 = T.tensor(0.).cuda()
weight_norm_set_2 = T.tensor(0.).cuda()
for name, w in model.named_parameters():
if 'noisy' not in name:
weight_norm_set_1 += w.norm().pow(2)
else:
weight_norm_set_2 += w.norm().pow(2)
# print(elbo.item())
normal_loss = nll*(1-flags.ELBO_weight) + elbo*flags.ELBO_weight + \
flags.IMAP_var_penalty_weight * imap_var
normal_loss = normal_loss + flags.l2 * weight_norm_set_1
penalty_weight = (flags.penalty_weight if
global_step >= flags.penalty_anneal_iters else 1.0)
normal_loss = normal_loss + penalty_weight * penalty
if penalty_weight > 1.0:
normal_loss = normal_loss / penalty_weight
pure_loss = nll.item() * (
1 - flags.ELBO_weight) + elbo.item() * flags.ELBO_weight
abnormal_loss = noisy_nll + flags.IMAP_var_penalty_weight * imap_var + flags.l2 * weight_norm_set_2
if train:
optimizer1.zero_grad()
normal_loss.backward(retain_graph=True)
optimizer1.step()
optimizer2.zero_grad()
abnormal_loss.backward()
optimizer2.step()
if iter % display_step == 0:
pretty_print(
np.int32(epoch),
np.int32(iter),
np.float32(pure_loss),
acc.detach().cpu().numpy(),
penalty.detach().cpu().numpy(),
)
total_loss += pure_loss
total_acc += acc.item()
total_penalty += penalty.item()
if train:
global_step += 1
iter += 1
if train:
global_imap = global_imap + imap.detach()
total_loss /= iter
total_acc /= iter
total_penalty /= iter
if train:
global_imap /= iter
return total_loss, total_acc, total_penalty
