def train_adversaries(trainloader,
valloader,
testloader,
writer,
navib_model=None):
if navib_model is None:
dummy_x, _, _ = trainset.__getitem__(0)
input_dim = dummy_x.size(0)
else:
navib_model.eval()
input_dim = args.latent_dim
output_dim = 5
hidden_dim = args.adv_hidden_dim
hidden_layers = 1
if navib_model is None:
name = 'Baseline_Adv'
else:
name = navib_model.name + '_New_Adv'
adv = models.Adv(name + str(hidden_layers),
input_dim=input_dim,
output_dim=output_dim,
hidden_dim=hidden_dim,
hidden_layers=hidden_layers).to(device)
opt = optim.Adam(adv.parameters(), lr=args.adv_lr, weight_decay=1e-4)
train_adv(adv, opt, trainloader, valloader, testloader, writer,
navib_model)
def train_adversaries(trainset, valset, testset, writer, navib_model=None):
if navib_model is None:
# For baseline, the adversaries train on the original data and not the latent space
dummy_x, _, _ = trainset.__getitem__(0)
input_dim = dummy_x.size(0)
else:
navib_model.eval()
input_dim = args.latent_dim # * 2
output_dim = 1
hidden_dim = args.adv_hidden_dim
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.adv_batch_size,
shuffle=True)
valloader = torch.utils.data.DataLoader(valset,
batch_size=args.adv_batch_size,
shuffle=False)
testloader = torch.utils.data.DataLoader(testset,
batch_size=args.adv_batch_size,
shuffle=False)
adv_hidden_layers = 3
if navib_model is None:
name = 'Baseline_Adv'
else:
name = navib_model.name + '_Adv'
adv = models.Adv(name + str(adv_hidden_layers),
input_dim=input_dim,
output_dim=output_dim,
hidden_dim=hidden_dim,
hidden_layers=adv_hidden_layers).to(device)
opt = optim.Adam(adv.parameters(), lr=args.adv_lr)
train_adv(adv, opt, trainloader, valloader, testloader, writer,
navib_model)
def train_navib(model, opt, trainloader, valloader, testloader, writer):
lr_scheduler = optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.65)
if args.comp_type == 'adv_training':
disc = models.Adv('Disc', input_dim=args.latent_dim, output_dim=2,
hidden_dim=args.adv_hidden_dim, hidden_layers=3).to(device)
disc_opt = optim.Adam(disc.parameters(), lr=args.disc_lr)
elif args.comp_type == 'uai':
disc = models.UAIDisentangler(latent_dim=args.latent_dim).to(device)
disc_opt = optim.Adam(disc.parameters(), lr=args.disc_lr)
for epoch in range(1, args.num_epochs + 1):
if args.comp_type == 'adv_training' or args.comp_type == 'uai':
navib_epoch(epoch, model, opt, trainloader, writer, tag='train',
disc=disc, disc_opt=disc_opt)
else:
navib_epoch(epoch, model, opt, trainloader, writer, tag='train')
navib_epoch(epoch, model, opt, valloader, writer, tag='val')
navib_epoch(epoch, model, opt, testloader, writer, tag='test')
if epoch % args.save_step == 0:
name = 'Navib_' + str(epoch)
path = os.path.join(model_path, name + '.pth')
model.name = name
torch.save(model.state_dict(), path)
lr_scheduler.step()
if args.alpha < args.alpha_max:
args.alpha *= args.alpha_gamma
name = 'Navib'
model.name = name
path = os.path.join(model_path, name + '.pth')
torch.save(model.state_dict(), path)
def start_adversarial_training(vae, writer, trainloader, valloader,
testloader):
num_hidden = 3
adv = models.Adv(name='Adv' + str(num_hidden),
input_dim=latent_dim,
output_dim=10,
hidden_layers=num_hidden).to(device)
adv_optimizer = optim.Adam(adv.parameters(), lr=args.adv_lr)
train_adversary(adv, vae, adv_optimizer, args.adv_num_epochs, writer,
trainloader, valloader, testloader)
def train_navib(model, opt, trainloader, valloader, testloader, writer):
lr_scheduler = optim.lr_scheduler.StepLR(opt, step_size=100, gamma=0.5)
best_val_acc = 0
if args.comp_type == 'adv_training':
disc = models.Adv('Disc',
input_dim=args.latent_dim,
output_dim=5,
hidden_dim=args.adv_hidden_dim,
hidden_layers=args.disc_hidden_layers).to(device)
disc_opt = optim.Adam(disc.parameters(), lr=args.disc_lr)
for epoch in range(1, args.num_epochs + 1):
if args.comp_type == 'adv_training':
navib_epoch(epoch,
model,
opt,
trainloader,
writer,
tag='train',
disc=disc,
disc_opt=disc_opt)
else:
navib_epoch(epoch, model, opt, trainloader, writer, tag='train')
val_acc = navib_epoch(epoch, model, opt, valloader, writer, tag='val')
test_Acc = navib_epoch(epoch,
model,
opt,
testloader,
writer,
tag='test')
if val_acc > best_val_acc:
name = 'Navib_best_val_acc'
model.name = name
path = os.path.join(model_path, name + '.pth')
torch.save(model.state_dict(), path)
best_val_acc = val_acc
message = 'Best val acc {}. test_acc{}. Saving model {}'.format(
best_val_acc, test_Acc, path)
print(message)
logf.write(message + '\n')
if epoch % args.save_step == 0:
name = 'Navib_ckpt_{}'.format(epoch)
model.name = name
path = os.path.join(model_path, name + '.pth')
torch.save(model.state_dict(), path)
print('Saving checkpoint {}'.format(path))
lr_scheduler.step()
if args.alpha < args.alpha_max:
args.alpha *= args.alpha_gamma
name = 'Navib'
model.name = name
path = os.path.join(model_path, name + '.pth')
torch.save(model.state_dict(), path)
def main(model, trainloader, valloader, testloader, opt, writer):
#lr_scheduler = optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.65)
if args.comp_type == 'adv_training':
num_hidden = 3
disc = models.Adv(name='Adv' + str(num_hidden),
input_dim=latent_dim,
output_dim=10,
hidden_layers=num_hidden).to(device)
disc_opt = optim.Adam(disc.parameters(), lr=args.disc_lr)
for epoch in range(args.num_epochs):
if args.comp_type == 'adv_training':
train_epoch(epoch, model, trainloader, opt, writer, disc, disc_opt)
else:
train_epoch(epoch, model, trainloader, opt, writer)
test_epoch(epoch, model, valloader, opt, writer, tag='val')
test_epoch(epoch, model, testloader, opt, writer, tag='test')
if args.comp_lambda < args.comp_lambda_max:
args.comp_lambda *= args.comp_gamma
#lr_scheduler.step()
print('Saving the model')
torch.save(model.state_dict(), net_path)