def train_valid_cycle(epochs=DEFAULT_NUM_EPOCHS, load_model=False):
# Gets the datasets
train_loader, valid_loader, test_loader = get_data_loader("binarized_mnist", DEFAULT_BATCH_SIZE)
# Creates VAE Model
model = VAE()
if load_model:
model.load_state_dict(torch.load('{}/model19.pt'.format(DEFAULT_SAVE_PATH)))
if torch.cuda.is_available():
model = model.cuda()
# Creates Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=DEFAULT_LR)
# Defines loss function
bce_loss = nn.BCELoss(size_average=False)
# Runs Train/Validation loop
for epoch in range(epochs):
train(epoch, model, optimizer, loss_fn=bce_loss, dataset=train_loader)
# Saves the model
torch.save(model.state_dict(), '{}/model_{}.pt'.format(DEFAULT_SAVE_PATH, epoch))
valid(epoch, model, loss_fn=bce_loss, dataset=valid_loader)
# Runs on Test Set
test(model, loss_fn=bce_loss, dataset=test_loader)
# Computes importance sampling logpx estimate on validation set
values = []
for index, data in enumerate(valid_loader):
data = data.cuda()
values.append(importance_sampling(model, data))
logpx = np.mean(np.concatenate(values))
print('log p(x) ~= {}'.format(logpx))
# Computes importance sampling logpx estimate on testing set
values = []
for index, data in enumerate(test_loader):
data = data.cuda()
values.append(importance_sampling(model, data))
logpx = np.mean(np.concatenate(values))
print('log p(x) ~= {}'.format(logpx))
'cuda': 1,
'hnf': 0
}
# q = Gaussian(hyper_config)
# q = Flow(hyper_config)
q = Flow1(hyper_config)
hyper_config['q'] = q
# Which gpu
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
print('Init model')
model = VAE(hyper_config)
if torch.cuda.is_available():
model.cuda()
print('\nModel:', hyper_config, '\n')
print('Load params for decoder')
path_to_load_variables = home + '/Documents/tmp/inference_suboptimality/vae_generator_3280.pt'
model.generator.load_state_dict(
torch.load(path_to_load_variables,
map_location=lambda storage, loc: storage))
print('loaded variables ' + path_to_load_variables)
print()
compute_local_opt = 0
compute_amort = 1
if compute_amort:
# 'encoder_arch': [[x_size,200],[200,200],[200,z_size*2]],
# 'decoder_arch': [[z_size,200],[200,200],[200,200],[200,200],[200,x_size]],
# 'q_dist': standard, #FFG_LN#,#hnf,#aux_nf,#flow1,#,
# 'cuda': 1
# }
# q = Gaussian(hyper_config)
# # q = Flow(hyper_config)
hyper_config['q'] = q(hyper_config)
print ('Init model')
model = VAE(hyper_config)
if torch.cuda.is_available():
model.cuda()
print('\nModel:', hyper_config,'\n')
# path_to_load_variables=''
# path_to_save_variables=home+'/Documents/tmp/inference_suboptimality/fashion_params/LE_binarized_fashion' #.pt'
# path_to_save_variables=home+'/Documents/tmp/inference_suboptimality/fashion_params/binarized_fashion_' #.pt'
# path_to_save_variables=home+'/Documents/tmp/pytorch_vae'+str(epochs)+'.pt'
# path_to_save_variables=this_dir+'/params_'+model_name+'_'
# path_to_save_variables=''