def train_cleitc(dataloader, seed, **kwargs):
"""
:param s_dataloaders:
:param t_dataloaders:
:param kwargs:
:return:
"""
autoencoder = AE(input_dim=kwargs['input_dim'],
latent_dim=kwargs['latent_dim'],
hidden_dims=kwargs['encoder_hidden_dims'],
dop=kwargs['dop']).to(kwargs['device'])
# get reference encoder
aux_ae = deepcopy(autoencoder)
aux_ae.encoder.load_state_dict(torch.load(os.path.join('./model_save/ae5000', f'ft_encoder_{seed}.pt')))
print('reference encoder loaded')
reference_encoder = aux_ae.encoder
# construct transmitter
transmitter = MLP(input_dim=kwargs['latent_dim'],
output_dim=kwargs['latent_dim'],
hidden_dims=[kwargs['latent_dim']]).to(kwargs['device'])
ae_eval_train_history = defaultdict(list)
ae_eval_test_history = defaultdict(list)
if kwargs['retrain_flag']:
cleit_params = [
autoencoder.parameters(),
transmitter.parameters()
]
cleit_optimizer = torch.optim.AdamW(chain(*cleit_params), lr=kwargs['lr'])
# start autoencoder pretraining
for epoch in range(int(kwargs['train_num_epochs'])):
if epoch % 1 == 0:
print(f'----Autoencoder Training Epoch {epoch} ----')
for step, batch in enumerate(dataloader):
ae_eval_train_history = cleit_train_step(ae=autoencoder,
reference_encoder=reference_encoder,
transmitter=transmitter,
batch=batch,
device=kwargs['device'],
optimizer=cleit_optimizer,
history=ae_eval_train_history)
torch.save(autoencoder.state_dict(), os.path.join(kwargs['model_save_folder'], 'cleit_ae.pt'))
torch.save(transmitter.state_dict(), os.path.join(kwargs['model_save_folder'], 'transmitter.pt'))
else:
try:
autoencoder.load_state_dict(torch.load(os.path.join(kwargs['model_save_folder'], 'cleit_ae.pt')))
transmitter.load_state_dict(torch.load(os.path.join(kwargs['model_save_folder'], 'transmitter.pt')))
except FileNotFoundError:
raise Exception("No pre-trained encoder")
encoder = EncoderDecoder(encoder=autoencoder.encoder,
decoder=transmitter).to(kwargs['device'])
return encoder, (ae_eval_train_history, ae_eval_test_history)
def train_ae(dataloader, **kwargs):
"""
:param s_dataloaders:
:param t_dataloaders:
:param kwargs:
:return:
"""
autoencoder = AE(input_dim=kwargs['input_dim'],
latent_dim=kwargs['latent_dim'],
hidden_dims=kwargs['encoder_hidden_dims'],
dop=kwargs['dop']).to(kwargs['device'])
ae_eval_train_history = defaultdict(list)
ae_eval_test_history = defaultdict(list)
if kwargs['retrain_flag']:
ae_optimizer = torch.optim.AdamW(autoencoder.parameters(),
lr=kwargs['lr'])
# start autoencoder pretraining
for epoch in range(int(kwargs['train_num_epochs'])):
if epoch % 50 == 0:
print(f'----Autoencoder Training Epoch {epoch} ----')
for step, batch in enumerate(dataloader):
ae_eval_train_history = ae_train_step(
ae=autoencoder,
batch=batch,
device=kwargs['device'],
optimizer=ae_optimizer,
history=ae_eval_train_history)
torch.save(autoencoder.state_dict(),
os.path.join(kwargs['model_save_folder'], 'ae.pt'))
else:
try:
autoencoder.load_state_dict(
torch.load(os.path.join(kwargs['model_save_folder'], 'ae.pt')))
except FileNotFoundError:
raise Exception("No pre-trained encoder")
return autoencoder.encoder, (ae_eval_train_history, ae_eval_test_history)
def train_cleita(dataloader, seed, **kwargs):
autoencoder = AE(input_dim=kwargs['input_dim'],
latent_dim=kwargs['latent_dim'],
hidden_dims=kwargs['encoder_hidden_dims'],
dop=kwargs['dop']).to(kwargs['device'])
# get reference encoder
aux_ae = deepcopy(autoencoder)
aux_ae.encoder.load_state_dict(
torch.load(os.path.join('./model_save', f'ft_encoder_{seed}.pt')))
print('reference encoder loaded')
reference_encoder = aux_ae.encoder
# construct transmitter
transmitter = MLP(input_dim=kwargs['latent_dim'],
output_dim=kwargs['latent_dim'],
hidden_dims=[kwargs['latent_dim']]).to(kwargs['device'])
confounding_classifier = MLP(input_dim=kwargs['latent_dim'],
output_dim=1,
hidden_dims=kwargs['classifier_hidden_dims'],
dop=kwargs['dop']).to(kwargs['device'])
ae_train_history = defaultdict(list)
ae_val_history = defaultdict(list)
critic_train_history = defaultdict(list)
gen_train_history = defaultdict(list)
if kwargs['retrain_flag']:
cleit_params = [autoencoder.parameters(), transmitter.parameters()]
cleit_optimizer = torch.optim.AdamW(chain(*cleit_params),
lr=kwargs['lr'])
classifier_optimizer = torch.optim.RMSprop(
confounding_classifier.parameters(), lr=kwargs['lr'])
for epoch in range(int(kwargs['train_num_epochs'])):
if epoch % 50 == 0:
print(f'confounder wgan training epoch {epoch}')
for step, batch in enumerate(dataloader):
critic_train_history = critic_train_step(
critic=confounding_classifier,
ae=autoencoder,
reference_encoder=reference_encoder,
transmitter=transmitter,
batch=batch,
device=kwargs['device'],
optimizer=classifier_optimizer,
history=critic_train_history,
# clip=0.1,
gp=10.0)
if (step + 1) % 5 == 0:
gen_train_history = gan_gen_train_step(
critic=confounding_classifier,
ae=autoencoder,
transmitter=transmitter,
batch=batch,
device=kwargs['device'],
optimizer=cleit_optimizer,
alpha=1.0,
history=gen_train_history)
torch.save(autoencoder.state_dict(),
os.path.join(kwargs['model_save_folder'], 'cleit_ae.pt'))
torch.save(transmitter.state_dict(),
os.path.join(kwargs['model_save_folder'], 'transmitter.pt'))
else:
try:
autoencoder.load_state_dict(
torch.load(
os.path.join(kwargs['model_save_folder'], 'cleit_ae.pt')))
transmitter.load_state_dict(
torch.load(
os.path.join(kwargs['model_save_folder'],
'transmitter.pt')))
except FileNotFoundError:
raise Exception("No pre-trained encoder")
encoder = EncoderDecoder(encoder=autoencoder.encoder,
decoder=transmitter).to(kwargs['device'])
return encoder, (ae_train_history, ae_val_history, critic_train_history,
gen_train_history)
label = label.to(device)
label = make_one_hot(label)
img = img.to(device)
_, output, hashed_layer = model(img)
if (i % 100 == 0 and epoch == 0) or (i % 500 == 0 and epoch > 0):
# PATH ASSUMES ONLY 1 STAGE
save_image(torch.cat((img, output)), "../results/ae_hash/images/train_check/{}_{}.jpg".format(epoch, i), nrow=batch_size)
loss, mse, msssim = criterion1(output, img)
loss_hash = criterion2(hashed_layer, label)
if torch.isnan(loss_hash).any():
torch.save(model.state_dict(), "nan_aya_wo_model_weights.pt")
torch.save(img, "nan_dene_wala_img_batch.pt")
assert not torch.isnan(loss_hash).any()
train_hash += loss_hash.cpu().item()
train_mse += mse.cpu().item()
train_msssim += msssim.cpu().item()
train_loss += loss.cpu().item()
for key, val in {'train_iter_loss': loss, 'train_iter_msssim': msssim, 'train_iter_hash': loss_hash, 'train_iter_mse': mse}.items():
writer.add_scalar(key, val.cpu().item(), train_iter_count)
train_iter_count += 1
total_loss = loss + loss_hash * cauchy_loss_weight
total_loss = loss
