Implementing multiple Keras Losses in PyTorch

vision
1

Hi,

I am trying to implement the following Keras codes in PyTorch.

This is a loss function I am trying to port to pytorch, I couldn’t seem to figure out how to implement this in Pytorch

def gradient_penalty_loss(y_true, y_pred, averaged_samples,
                          gradient_penalty_weight):
    
    gradients = K.gradients(y_pred, averaged_samples)[0]
    gradients_sqr = K.square(gradients)
    gradients_sqr_sum = K.sum(gradients_sqr,
                              axis=np.arange(1, len(gradients_sqr.shape)))
    gradient_l2_norm = K.sqrt(gradients_sqr_sum)
    gradient_penalty = gradient_penalty_weight * K.square(1 - gradient_l2_norm)
    return K.mean(gradient_penalty)

  partial_gp_loss = partial(gradient_penalty_loss,
                          averaged_samples=averaged_samples,
                          gradient_penalty_weight=GRADIENT_PENALTY_WEIGHT)

Any help regarding implementing the loss function and porting the above codes to pytorch will be a great help.

UPDATE: I found the pytorch implementation in this link. I am posting it here so that its going to help someone out in the future.