gninatorch.losses module

class gninatorch.losses.AffinityLoss(reduction: str = 'mean', delta: float = 1.0, penalty: float = 0.0, pseudo_huber: bool = False, scale: float = 1.0)[source]

Bases: Module

GNINA affinity loss.

Parameters:
  • reduction (str) – Reduction method (mean or sum)

  • delta (float) – Scaling factor

  • penalty (float) – Penalty factor

  • pseudo_huber (bool) – Use pseudo-huber loss as opposed to L2 loss

  • scale (float) – Scaling factor for the loss

Notes

Translated from the original custom Caffe layer. Not all functionality is implemented.

https://github.com/gnina/gnina/blob/master/caffe/src/caffe/layers/affinity_loss_layer.cpp

The scale parameter is different from the original implementation. In the original Caffe implementation, the scale parameter is used to scale the gradients in the backward pass. Here the scale parameter scales the loss function directly in the forward pass.

Definition of pseudo-Huber loss: https://en.wikipedia.org/wiki/Huber_loss#Pseudo-Huber_loss_function

forward(input: Tensor, target: Tensor) Tensor[source]
Parameters:
  • input (Tensor) – Predicted values

  • target (Tensor) – Target values

Returns:

Loss

Return type:

torch.Tensor

Notes

Binding affinity (pK) is positive for good poses and negative for bad poses (and zero if unknown). This allows to distinguish good poses from bad poses (to which a penalty is applied) without explicitly using the labels or the RMSD.

class gninatorch.losses.ScaledNLLLoss(scale: float = 1.0, reduction: str = 'mean')[source]

Bases: Module

Scaled NLLLoss.

Parameters:
  • scale (float) – Scaling factor for the loss

  • reduction (str) – Reduction method (mean or sum)

forward(input: Tensor, target: Tensor) Tensor[source]
Parameters:
  • input (Tensor) – Predicted values

  • target (Tensor) – Target values

Returns:

Loss

Return type:

torch.Tensor