Need help for custom loss backward() function

How should I implement my very own custom loss function.
There are two ways to implement using a custom loss class and writing function with init and forward.
I had the code already written for TensorFlow in class, therefore I chose to go along.
For PyTorch, I tried explored the below relevant discussions, but could not figure how to design the backwards function for same.

Custom Loss:

class MSE_ITL_RegLoss(torch.nn.Module):
    def __init__(self, latent_code, prior, regularization_factor, kernel, batch_size):
        super (MSE_ITL_RegLoss).__init__()
        
        self.latent_code = latent_code         
        self.regularization_factor = regularization_factor
        self.prior = prior
        self.kernel = kernel
        self.batch_size = batch_size
        
    
    def evaluate_ITL(self, first_pdf, second_pdf):       
        error = first_pdf - second_pdf
        pairwise_error = torch.reshape(torch.unsqueeze(error ,1) - error, (-1,))
        temp = torch.exp(-(np.sqrt(2) * self.kernel) * torch.square(pairwise_error_t))
        return torch.sum(temp)/self.batch_size**2
        
    
    def cs_div (self):
        p_x = self.evaluate_ITL(self.latent_code, self.latent_code)
        q_x = self.evaluate_ITL(self.prior,self.prior)
        p_x__q_x = self.evaluate_ITL(self.latent_code, self.prior) 
        return p_x__q_x / torch.sqrt(p_x * q_x)
        
        
    def forward(self, y_true, y_pred):       
        mse = torch.mean(torch.square(y_true_t-y_pred_t))
        cs_div = self.cs_div()       
        final_loss = mse + self.regularization_factor * cs_div
        return final_loss

Model:

class modelMSE_CSD (nn.Module):
    def __init__(self, input_size, bottleneck_dim):
        super(modelMSE_CSD, self).__init__()

        self.bottleneck_dim = bottleneck_dim
        self.input_size = input_size
        
        self.encoder = nn.Sequential(
            nn.Linear(self.input_size, 800),
            nn.ReLU(),
            nn.Linear(800,400),
            nn.ReLU(),
            nn.Linear(400, self.bottleneck_dim)
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(self.bottleneck_dim, 400),
            nn.ReLU(),
            nn.Linear(400,800),
            nn.ReLU(),
            nn.Linear(800, self.input_size),
            nn.Sigmoid()
        )
        
    def forward(self,x):
        e_out = self.encoder(x)
        d_out = self.decoder(e_out)
        return e_out, d_out

Training Loop:

losses = []
for epoch in range(EPOCHS):
    running_loss = 0.0
    
    for batch_num, (data, target) in enumerate(dataloader):
        
        data = Variable(data, requires_grad = True)
        target = Variable(target, requires_grad = True)

        model_csd.train()
        latent_code, decoded_layer = model_csd(data)
        
        loss = MSE_ITL_RegLoss(latent_code=latent_code, prior=PRIOR_DIST, 
                               regularization_factor=REGULARIZER, kernel=KERNEL, batch_size=BATCH_SIZE)

        optimizer.zero_grad()
        **loss.backward()**
        optimizer.step()
        
        running_loss += loss.item() * data.size[0]
    losses.append(running_loss/len(dataset))
    print("Epoch: ",epoch,"Loss :", running_loss/len(dataset))```



Please suggest if the other method of declaring the custom loss is more easier.
Any links/resources would be appreciated.

You don’t need to define backward method if your custom loss class is inherited from nn.Module, it’s already there.

Hi @my3bikaht , thank you for your response.

While running the training loop (code mentioned above), I get the error that AttributeError: 'MSE_ITL_RegLoss' object has no attribute 'backward'

This is the error traceback:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-13-fc852af04b41> in <module>
     15 
     16         optimizer.zero_grad()
---> 17         loss.backward()
     18         optimizer.step()
     19 

c:\...\torch\nn\modules\module.py in __getattr__(self, name)
   1176                 return modules[name]
   1177         raise AttributeError("'{}' object has no attribute '{}'".format(
-> 1178             type(self).__name__, name))
   1179 
   1180     def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -> None:

AttributeError: 'MSE_ITL_RegLoss' object has no attribute 'backward'

Is there anything that can be done?

You are creating an object called loss from the MSE_ITL_RegLoss class:

loss = MSE_ITL_RegLoss(latent_code=latent_code, prior=PRIOR_DIST, 
                               regularization_factor=REGULARIZER, kernel=KERNEL, batch_size=BATCH_SIZE)

but are never actually calculating the final_loss computed in the forward.
Based on the class definition you would need to call this object (same as calling a model with inputs) to get the loss:

loss = MSE_ITL_RegLoss(latent_code=latent_code, prior=PRIOR_DIST, 
                               regularization_factor=REGULARIZER, kernel=KERNEL, batch_size=BATCH_SIZE)
final_loss = loss(y_true, y_pred)
final_loss.backward()