Custom loss function using non-pytorch functions

username: lbry
cooked: <p>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.<br>
My goal is to make the generated rules not random but instead fit to a certain set of chemical reactions.</p>
<p>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.</p>
<p>Here is my main class and training loop</p>
<pre><code class="lang-auto">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()
</code></pre>
<p>Here is the function I would ideally like to use as some sort of loss function.</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>I am rather new to pytorch so let me know if this is not really possible!</p>