Loss function for Bayesian Neural Network: backward() issues

username: domage
cooked: <p>Hello PyTorch community,</p>
<p>I’m currently working on a Bayesian Neural Network (BNN) and I’ve encountered an issue with my custom loss function. I have a BnnLoss class that calculates the base loss and adds a complexity cost (KL divergence) for BNNs. Here’s the code for the class:</p>
<pre><code class="lang-auto">class BnnLoss(torch.nn.Module):
    def __init__(self, base_loss: torch.nn.Module, kl_weight: float, model: torch.nn.Module = None):
        super().__init__()
        self.base_loss = base_loss
        self.kl_weight = kl_weight
        self.model = model if isinstance(model, BnnFourierNet) else None

    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -&gt; torch.Tensor:
        data_loss = self.base_loss(outputs.squeeze(), targets)

        if self.model is not None:
            complexity_cost = self.model.calculate_kl_divergence()
            n_data = targets.shape[0]
            return data_loss + (self.kl_weight/n_data) * complexity_cost
        else:        
            return data_loss
</code></pre>
<p>When I use this class in my training loop and call <code>total_loss.backward()</code> , the training result is not as expected. The standard deviation of the model’s output becomes very low, thus acting as a base_loss function, without any modifications.</p>
<p>However, if I calculate the loss directly in the training loop without using the BnnLoss class, everything works as expected. Here is the code of the training cycle, where I try to compare both approaches:</p>
<pre><code class="lang-auto">def train_solo_cycle(model, loss, device, conf, trainloader, opt, prof=None):
    for ep in range(conf["epochs"]):
        start_time = time.time()
        for batch in trainloader:
            opt.zero_grad()
            inputs = batch[0].to(device)
            targets = batch[1].to(device)
            outputs = model.forward(inputs)

            # That loss calculation is implemented directly in the trainig cycle
            # I use base_loss (e.g. torch.nn.BCELoss()) here directly. Works as intended
            if isinstance(model, BnnFourierNet):
                data_loss = loss.base_loss(outputs.squeeze(), targets)
                complexity_cost = model.calculate_kl_divergence()
                n_data = targets.shape[0] 
                total_loss_solo_cycle = data_loss + conf["kl_weight"]/n_data * complexity_cost
            # Here, i calculate the loss using BnnLoss class       
            total_loss = loss(outputs.squeeze(), targets)
            # That yelds true      
            torch.eq(total_loss, total_loss_solo_cycle)
            # If I use that backward() call from the BnnLoss class, the training results would be the same              
            # as if I use torch.nn.BCELoss() without any modifications (deviation of the model's output becomes very low)
            total_loss.backward()

            # If I use that backward() call, from the variable directry in the training cycle
            # training results would be as expected
            #total_loss_solo_cycle.backward()

            opt.step()  
    return model
</code></pre>
<p>What’s strange is that when I compare the tensors generated by the BnnLoss class and in the training loop, they are equal. However, the behavior of the backward operation seems to differ depending on whether it’s called from within the BnnLoss class or directly in the training loop.</p>
<p>I’ve tried using <code>torch.autograd.set_detect_anomaly(True)</code> to identify the source of the issue, but it didn’t provide any additional insights.</p>
<p>I would appreciate any insights or suggestions on what might be causing this issue and how to resolve it.</p>
<p>Thank you in advance for your help!</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="domage" data-post="1" data-topic="180712">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/domage/48/60779_2.png" class="avatar"> domage:</div>
<blockquote>
<p>The standard deviation of the model’s output becomes very low, thus acting as a base_loss function, without any modifications.</p>
</blockquote>
</aside>
<p>Did you verify that your custom <code>BnnLoss</code> method is indeed using the <code>if self.model is not None</code> path?</p>

username: domage
cooked: <p>Yes. I checked that in the debug mode. Also, changing the forward function to the following implementation did not influence on the final result:</p>
<pre><code class="lang-auto">def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -&gt; torch.Tensor:
        complexity_cost = self.model.calculate_kl_divergence()
        n_data = targets.shape[0] 
        return self.base_loss(outputs.squeeze(), targets) + (self.kl_weight/n_data) * complexity_cost
</code></pre>

username: domage
cooked: <p>Ok. I found the error. I initialize the BnnLoss with the empty model, but I implement several training sessions, using the copy.deepcopy(base_model) in the other part of the code. So, the BnnLoss never gets the updated status of the model, working constantly on the unprocessed one.</p>