def test_gradient_penalty_non_zero(self):
# Test to verify that a non-zero gradient penalty is computed on the from the first training step
with TemporaryDirectory() as tmpdirname:
latent = np.random.rand(64, 1, 512)
os.makedirs(os.path.dirname(tmpdirname + '/encoded_smiles.latent'),
exist_ok=True)
with open(tmpdirname + '/encoded_smiles.latent', 'w') as f:
json.dump(latent.tolist(), f)
C = CreateModelRunner(input_data_path=tmpdirname +
'/encoded_smiles.latent',
output_model_folder=tmpdirname)
C.run()
D = Discriminator.load(tmpdirname + '/discriminator.txt')
G = Generator.load(tmpdirname + '/generator.txt')
json_smiles = open(tmpdirname + '/encoded_smiles.latent', "r")
latent_space_mols = np.array(json.load(json_smiles))
testSampler = Sampler(G)
latent_space_mols = latent_space_mols.reshape(
latent_space_mols.shape[0], 512)
T = torch.cuda.FloatTensor
G.cuda()
D.cuda()
dataloader = torch.utils.data.DataLoader(
LatentMolsDataset(latent_space_mols),
shuffle=True,
batch_size=64,
drop_last=True)
for _, real_mols in enumerate(dataloader):
real_mols = real_mols.type(T)
fake_mols = testSampler.sample(real_mols.shape[0])
alpha = T(np.random.random((real_mols.size(0), 1)))
interpolates = (alpha * real_mols +
((1 - alpha) * fake_mols)).requires_grad_(True)
d_interpolates = D(interpolates)
fake = T(real_mols.shape[0], 1).fill_(1.0)
gradients = autograd.grad(
outputs=d_interpolates,
inputs=interpolates,
grad_outputs=fake,
create_graph=True,
retain_graph=True,
only_inputs=True,
)[0]
gradients = gradients.view(gradients.size(0), -1)
gradient_penalty = ((gradients.norm(2, dim=1) - 1)**2).mean()
self.assertTrue(gradient_penalty.data != 0)
break
def test_discriminator_shape(self):
# Test to verify that the same dimension network is created invariant of smiles input file size
with TemporaryDirectory() as tmpdirname:
for j in [1, 64, 256, 1024]:
latent = np.random.rand(j, 1, 512)
os.makedirs(os.path.dirname(tmpdirname +
'/encoded_smiles.latent'),
exist_ok=True)
with open(tmpdirname + '/encoded_smiles.latent', 'w') as f:
json.dump(latent.tolist(), f)
C = CreateModelRunner(input_data_path=tmpdirname +
'/encoded_smiles.latent',
output_model_folder=tmpdirname)
C.run()
D = Discriminator.load(tmpdirname + '/discriminator.txt')
D_params = []
for param in D.parameters():
D_params.append(param.view(-1))
D_params = torch.cat(D_params)
reference = 394241
self.assertEqual(D_params.shape[0], reference,
"Network does not match expected size")
def test_separate_optimizers(self):
# Verify that two different instances of the optimizer is created using the TrainModelRunner.py initialization
# This ensures the two components train separately
with TemporaryDirectory() as tmpdirname:
latent = np.random.rand(64, 1, 512)
os.makedirs(os.path.dirname(tmpdirname + '/encoded_smiles.latent'),
exist_ok=True)
with open(tmpdirname + '/encoded_smiles.latent', 'w') as f:
json.dump(latent.tolist(), f)
C = CreateModelRunner(input_data_path=tmpdirname +
'/encoded_smiles.latent',
output_model_folder=tmpdirname)
C.run()
D = Discriminator.load(tmpdirname + '/discriminator.txt')
G = Generator.load(tmpdirname + '/generator.txt')
optimizer_G = torch.optim.Adam(G.parameters())
optimizer_D = torch.optim.Adam(D.parameters())
self.assertTrue(type(optimizer_G) == type(
optimizer_D)) # must return the same type of object
self.assertTrue(
optimizer_G
is not optimizer_D) # object identity MUST be different
def __init__(self,
input_data_path,
output_model_folder,
decode_mols_save_path='',
n_epochs=200,
starting_epoch=1,
batch_size=64,
lr=0.0002,
b1=0.5,
b2=0.999,
n_critic=5,
sample_interval=10,
save_interval=100,
sample_after_training=100,
message=""):
self.message = message
# init params
self.input_data_path = input_data_path
self.output_model_folder = output_model_folder
self.n_epochs = n_epochs
self.starting_epoch = starting_epoch
self.batch_size = batch_size
self.lr = lr
self.b1 = b1
self.b2 = b2
self.n_critic = n_critic
self.sample_interval = sample_interval
self.save_interval = save_interval
self.sample_after_training = sample_after_training
self.decode_mols_save_path = decode_mols_save_path
# initialize dataloader
json_smiles = open(self.input_data_path, "r")
latent_space_mols = np.array(json.load(json_smiles))
latent_space_mols = latent_space_mols.reshape(
latent_space_mols.shape[0], 512)
self.dataloader = torch.utils.data.DataLoader(
LatentMolsDataset(latent_space_mols),
shuffle=True,
batch_size=self.batch_size)
# load discriminator
discriminator_name = 'discriminator.txt' if self.starting_epoch == 1 else str(
self.starting_epoch - 1) + '_discriminator.txt'
discriminator_path = os.path.join(output_model_folder,
discriminator_name)
self.D = Discriminator.load(discriminator_path)
# load generator
generator_name = 'generator.txt' if self.starting_epoch == 1 else str(
self.starting_epoch - 1) + '_generator.txt'
generator_path = os.path.join(output_model_folder, generator_name)
self.G = Generator.load(generator_path)
# initialize sampler
self.Sampler = Sampler(self.G)
# initialize optimizer
self.optimizer_G = torch.optim.Adam(self.G.parameters(),
lr=self.lr,
betas=(self.b1, self.b2))
self.optimizer_D = torch.optim.Adam(self.D.parameters(),
lr=self.lr,
betas=(self.b1, self.b2))
# Tensor
cuda = True if torch.cuda.is_available() else False
if cuda:
self.G.cuda()
self.D.cuda()
self.Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
def test_model_trains(self):
# Performs one step of training and verifies that the weights are updated, implying some training occurs.
with TemporaryDirectory() as tmpdirname:
T = torch.cuda.FloatTensor
latent = np.random.rand(64, 1, 512)
os.makedirs(os.path.dirname(tmpdirname + '/encoded_smiles.latent'),
exist_ok=True)
with open(tmpdirname + '/encoded_smiles.latent', 'w') as f:
json.dump(latent.tolist(), f)
C = CreateModelRunner(input_data_path=tmpdirname +
'/encoded_smiles.latent',
output_model_folder=tmpdirname)
C.run()
D = Discriminator.load(tmpdirname + '/discriminator.txt')
G = Generator.load(tmpdirname + '/generator.txt')
G.cuda()
D.cuda()
optimizer_G = torch.optim.Adam(G.parameters())
optimizer_D = torch.optim.Adam(D.parameters())
json_smiles = open(tmpdirname + '/encoded_smiles.latent', "r")
latent_space_mols = np.array(json.load(json_smiles))
testSampler = Sampler(G)
latent_space_mols = latent_space_mols.reshape(
latent_space_mols.shape[0], 512)
dataloader = torch.utils.data.DataLoader(
LatentMolsDataset(latent_space_mols),
shuffle=True,
batch_size=64,
drop_last=True)
for _, real_mols in enumerate(dataloader):
real_mols = real_mols.type(T)
before_G_params = []
before_D_params = []
for param in G.parameters():
before_G_params.append(param.view(-1))
before_G_params = torch.cat(before_G_params)
for param in D.parameters():
before_D_params.append(param.view(-1))
before_D_params = torch.cat(before_D_params)
optimizer_D.zero_grad()
fake_mols = testSampler.sample(real_mols.shape[0])
real_validity = D(real_mols)
fake_validity = D(fake_mols)
#It is not relevant to compute gradient penalty. The test is only interested in if there is a change in
#the weights (training), not in giving proper training
d_loss = -torch.mean(real_validity) + torch.mean(fake_validity)
d_loss.backward()
optimizer_D.step()
optimizer_G.zero_grad()
fake_mols = testSampler.sample(real_mols.shape[0])
fake_validity = D(fake_mols)
g_loss = -torch.mean(fake_validity)
g_loss.backward()
optimizer_G.step()
after_G_params = []
after_D_params = []
for param in G.parameters():
after_G_params.append(param.view(-1))
after_G_params = torch.cat(after_G_params)
for param in D.parameters():
after_D_params.append(param.view(-1))
after_D_params = torch.cat(after_D_params)
self.assertTrue(
torch.any(torch.ne(after_G_params, before_G_params)))
self.assertTrue(
torch.any(torch.ne(after_D_params, before_D_params)))
break
