I have a generator which I am using to create currently random chemical rules. The rules are generated by inputting into the latent space of an autoencoder with random numbers.
My goal is to make the generated rules not random but instead fit to a certain set of chemical reactions.
My idea is that I train a layer that can output values that when input to the latent space produce a rule that adequately explains a set of similar chemical reactions. To do this I have created a function that outputs the score of a given rule (see code below). I would like to use it as a loss function if possible. Currently I pass it to a mseloss but that errors as the tensors don’t have a grad_fn.
Here is my main class and training loop
class RuleLayer(nn.Module):
def __init__(self):
super().__init__()
self.layer = nn.Linear(1, 16)
def forward(self, x):
return self.layer(x)
current_cluster = list(grouped[group]) # group of similar reactions that need a rule
data = [[0.0]] * 50
data = torch.tensor(data)
loader = torch.utils.data.DataLoader(dataset=data, batch_size=16, shuffle=True)
model = RuleLayer()
loss_function = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.005)
epochs = 50
for epoch in range(epochs):
print("Epoch {} of {}".format(epoch + 1, epochs))
for batch in loader:
optimizer.zero_grad()
output = model(batch)
batch_smarts = []
for value in output: # For each instance in the batch generate a rule
latent = x_latent[0:1]
value = value.detach().numpy()
latent_r = latent + value
batch_smarts.append(latent_to_smiles(latent_r, embed))
losses = []
for batch_smart in batch_smarts: # Calculate the score of each generated rule
losses.append(rule_accuracy_loss(batch_smart))
losses = torch.tensor(losses)
comparison = torch.tensor([1.0] * len(batch)) # The ideal rule would explain all the reactions perfectly so the comparison for the mseloss should be 1.0
loss = loss_function(losses, comparison)
loss.backward()
optimizer.step()
Here is the function I would ideally like to use as some sort of loss function.
def rule_accuracy_loss(new_smart):
# Check if rule generated is actually a valid rule if not return 0
try:
rd_reaction = AllChem.ReactionFromSmarts(new_smart)
except ValueError:
return 0.0
correct_count = 0
# Calculate the amount of reactions the rule works with
for reaction in current_cluster.keys():
original_products, reactants = get_reactant_product(reaction, current_cluster)
success, products = works_with_rule(rd_reaction, reactants)
if not success:
continue
correct = products_equal(products, original_products)
if correct[1]:
correct_count += 1
total_reactions = len(current_cluster.keys())
# Return value above 0 (as rule is syntactically valid) plus rule score
return 0.1 + ((correct_count / total_reactions) * 0.9)
I am rather new to pytorch so let me know if this is not really possible!