def objective(params):
"""
Objective function to be minimized: loss with respect to our hyperparameters.
"""
enc_kernel1 = int(params[0])
enc_kernel2 = int(params[1])
enc_kernel3 = int(params[2])
dec_kernel1 = int(params[3])
dec_kernel2 = int(params[4])
dec_kernel3 = int(params[5])
# Contact matrices are 21x21
input_dim = 441
encoder = Encoder(input_size=input_dim,
latent_size=8,
kernel1=enc_kernel1,
kernel2=enc_kernel2,
kernel3=enc_kernel3)
decoder = Decoder(latent_dim=8,
output_size=input_size,
kernel1=dec_kernel1,
kernel2=dec_kernel2,
kernel3=dec_kernel3)
vae = VAE(encoder, decoder)
criterion = nn.MSELoss()
use_cuda = args.use_cuda
if use_cuda:
encoder = encoder.cuda()
deconder = decoder.cuda()
vae = vae.cuda()
criterion = criterion.cuda()
optimizer = optim.Adam(vae.parameters(), lr=0.0001)
epoch_loss = 0
total_loss = 0
for epoch in range(100):
for i, data in enumerate(trainloader, 0):
inputs = data['cont_matrix']
inputs = inputs.resize_(args.batch_size, 1, 21, 21)
inputs = inputs.float()
if use_cuda:
inputs = inputs.cuda()
inputs = Variable(inputs)
optimizer.zero_grad()
dec = vae(inputs)
ll = latent_loss(vae.z_mean, vae.z_sigma)
loss = criterion(dec, inputs) + ll
loss.backward()
optimizer.step()
epoch_loss = loss.data[0]
print(epoch, epoch_loss)
total_loss += epoch_loss
return total_loss
def main():
use_cuda = args.use_cuda
train_data = UnlabeledContact(data=args.data_dir)
print('Number of samples: {}'.format(len(train_data)))
trainloader = DataLoader(train_data, batch_size=args.batch_size)
# Contact matrices are 21x21
input_size = 441
encoder = Encoder(input_size=input_size, latent_size=3)
decoder = Decoder(latent_size=3, output_size=input_size)
vae = VAE(encoder, decoder, use_cuda=use_cuda)
criterion = nn.MSELoss()
if use_cuda:
encoder = nn.DataParallel(encoder)
decoder = nn.DataParallel(decoder)
encoder = encoder.cuda().half()
decoder = decoder.cuda().half()
vae = nn.DataParallel(vae)
vae = vae.cuda().half()
criterion = criterion.cuda().half()
optimizer = optim.SGD(vae.parameters(), lr=0.01)
clock = AverageMeter(name='clock16', rank=0)
epoch_loss = 0
total_loss = 0
end = time.time()
for epoch in range(15):
for batch_idx, data in enumerate(trainloader):
inputs = data['cont_matrix']
# inputs = inputs.resize_(args.batch_size, 1, 21, 21)
inputs = inputs.float()
if use_cuda:
inputs = inputs.cuda().half()
inputs = Variable(inputs)
optimizer.zero_grad()
dec = vae(inputs)
ll = latent_loss(vae.z_mean, vae.z_sigma)
loss = criterion(dec, inputs) + ll
loss.backward()
optimizer.step()
epoch_loss += loss.data[0]
clock.update(time.time() - end)
end = time.time()
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(trainloader.dataset),
100. * batch_idx / len(trainloader), loss.data[0]))
clock.save(path='/home/ygx/libraries/mds/molecules/molecules/linear_vae')
def test(x, vae, vqvae_model, do_use_cuda=False, save_img_path=None):
x_d = vae(x.contiguous().view(x.shape[0], -1))
x_di = x_d.contiguous().view(x_d.shape[0], probs_size, dsize, dsize)
xi = x.contiguous().view(x.shape[0], 1, dsize, dsize)
dmll_loss = discretized_mix_logistic_loss(x_di,
2 * xi - 1,
nr_mix=nr_mix,
use_cuda=do_use_cuda)
kl_loss = kl_weight * latent_loss(vae.z_mean, vae.z_sigma)
test_loss = dmll_loss + kl_loss
return test_loss
def train(args):
dataset = args.dataset
if dataset == "yaleb":
input_dim = 32256
batch_size = 32
z_dim = 100
hidden = 300
criterion = nn.BCELoss()
max_epochs = 100
if dataset == "adult":
input_dim = 108
batch_size = 64
z_dim = 2
hidden = 100
criterion = nn.MSELoss()
max_epochs = 30
if dataset == "german":
input_dim = 61
batch_size = 64
z_dim = 2
hidden = 100
criterion = nn.MSELoss()
max_epochs = 30
dataloader, dataloader_test = load_data(dataset, batch_size)
encoder = Encoder(input_dim, hidden, hidden)
decoder = Decoder(z_dim, hidden, input_dim)
vae = VAE(encoder, decoder, hidden, z_dim)
optimizer = optim.Adam(vae.parameters(), lr=0.00001)
l = None
ls = []
for epoch in range(max_epochs):
for i, data in enumerate(dataloader, 0):
inputs, classes, _ = data
optimizer.zero_grad()
dec = vae(inputs)
ll = latent_loss(vae.z_mean, vae.z_sigma)
loss = criterion(dec, inputs) + ll
loss.backward()
optimizer.step()
l = loss.data.item()
ls.append(l)
print(epoch, l)
with open(PROJECT_DIR / f"models/{dataset}_vae", "wb") as f:
torch.save(vae, f)
plt.plot(ls)
plt.show()
def train(epoch, model, optimizer, train_loader, do_checkpoint, do_use_cuda):
latent_losses = []
dmll_losses = []
kl_weight = min(1.0, epoch * 1e-2 + .1)
for batch_idx, (data, _) in enumerate(train_loader):
start_time = time.time()
if do_use_cuda:
x = Variable(data, requires_grad=False).cuda()
else:
x = Variable(data, requires_grad=False)
optimizer.zero_grad()
x_d = vae(x.contiguous().view(x.shape[0], -1))
x_di = x_d.contiguous().view(x_d.shape[0], probs_size, dsize, dsize)
dmll_loss = discretized_mix_logistic_loss(x_di,
2 * x - 1,
nr_mix=nr_mix,
use_cuda=do_use_cuda)
#x_di = x_d.contiguous().view(x_d.shape[0], 1, dsize, dsize)
#dmll_loss = mse_loss(x_di, x)
kl_loss = kl_weight * latent_loss(vae.z_mean, vae.z_sigma)
loss = dmll_loss + kl_loss
loss.backward()
optimizer.step()
latent_losses.append(kl_loss.cpu().data)
dmll_losses.append(dmll_loss.cpu().data)
if not batch_idx % 500:
print 'Train Epoch: {} [{}/{} ({:.0f}%)]\tKL Loss: {} MSE Loss: {} Time: {}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / float(len(train_loader)),
np.asarray(latent_losses).mean(0),
np.asarray(dmll_losses).mean(0),
time.time() - start_time)
state = {
'epoch': epoch,
'state_dict': vae.state_dict(),
'dmll_losses': np.asarray(dmll_losses).mean(0),
'latent_losses': np.asarray(latent_losses).mean(0),
'optimizer': optimizer.state_dict(),
}
return model, optimizer, state
def train(epoch, model, optimizer, train_loader, do_checkpoint, do_use_cuda):
latent_losses = []
mse_losses = []
kl_weight = min(1.0, epoch * 1e-2)
for batch_idx, (data, _) in enumerate(train_loader):
start_time = time.time()
if do_use_cuda:
x = Variable(data, requires_grad=False).cuda()
else:
x = Variable(data, requires_grad=False)
optimizer.zero_grad()
dec = vae(x)
kl = kl_weight * latent_loss(vae.z_mean, vae.z_sigma)
mse_loss = criterion(dec, x)
loss = mse_loss + kl
loss.backward()
optimizer.step()
latent_losses.append(kl.cpu().data)
mse_losses.append(mse_loss.cpu().data)
if not batch_idx % 500:
print 'Train Epoch: {} [{}/{} ({:.0f}%)]\tKL Loss: {} MSE Loss: {} Time: {}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / float(len(train_loader)),
np.asarray(latent_losses).mean(0),
np.asarray(mse_losses).mean(0),
time.time() - start_time)
state = {
'epoch': epoch,
'state_dict': vae.state_dict(),
'mse_losses': np.asarray(mse_losses).mean(0),
'latent_losses': np.asarray(latent_losses).mean(0),
'optimizer': optimizer.state_dict(),
}
return model, optimizer, state
def main():
use_cuda = args.use_cuda
train_data = UnlabeledContact(data=args.data_dir)
print('Number of samples: {}'.format(len(train_data)))
trainloader = DataLoader(train_data, batch_size=args.batch_size)
# Contact matrices are 21x21
input_size = 441
encoder = Encoder(input_size=input_size, latent_size=3)
decoder = Decoder(latent_size=3, output_size=input_size)
vae = VAE(encoder, decoder, use_cuda=use_cuda)
criterion = nn.MSELoss()
if use_cuda:
encoder = encoder.cuda().half()
decoder = decoder.cuda().half()
vae = vae.cuda().half()
criterion = criterion.cuda().half()
optimizer = optim.SGD(vae.parameters(), lr=0.01)
epoch_loss = 0
total_loss = 0
for epoch in range(100):
for batch_idx, data in enumerate(trainloader):
inputs = data['cont_matrix']
inputs = inputs.resize_(args.batch_size, 1, 21, 21)
inputs = inputs.float()
if use_cuda:
inputs = inputs.cuda().half()
inputs = Variable(inputs)
optimizer.zero_grad()
dec = vae(inputs)
ll = latent_loss(vae.z_mean, vae.z_sigma)
loss = criterion(dec, inputs) + ll
loss.backward()
optimizer.step()
epoch_loss += loss.data[0]
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(trainloader.dataset),
100. * batch_idx / len(trainloader), loss.data[0]))
if epoch < 10:
# Get latent encoding
latent_array = encoder(inputs).data[0].cpu().float().numpy()
filename = 'latent_epoch' + str(epoch)
np.save('./latent_saves/' + filename, latent_array)
# Get reconstructed image
reconstructed_array = vae(
inputs).data[0].cpu().float().numpy().reshape(21, 21)
recon_filename = 'reconstructed_epoch' + str(epoch)
np.save('./reconstruct_saves/' + recon_filename,
reconstructed_array)
if epoch % 10 == 0:
torch.save(vae.state_dict(), args.save_path + 'epoch' + str(epoch))
latent_array = encoder(inputs).data[0].cpu().float().numpy()
filename = 'latent_epoch' + str(epoch)
np.save('./latent_saves/' + filename, latent_array)
reconstructed_array = vae(
inputs).data[0].cpu().float().numpy().reshape(21, 21)
recon_filename = 'reconstructed_epoch' + str(epoch)
np.save('./reconstruct_saves/' + recon_filename,
reconstructed_array)
def test(x, vae, vqvae_model, do_use_cuda=False, save_img_path=None):
dec = vae(x)
kl = latent_loss(vae.z_mean, vae.z_sigma)
loss = criterion(dec, x) + kl
test_loss = loss.cpu().data.mean()
return test_loss