multiTrain.py

 
 
 
from torch import nn
 
 
class MultiTrainLoss(nn.Module):
    """
    Sum of cross-entropies for training against LCAS and mass hypotheses.
 
    Args:
        alpha_mass (float): Weight of mass cross-entropy term in the loss.
        ignore_index (int): Index to ignore in the computation (e.g. padding).
        reduction (str): Type of reduction to be applied on the batch (``sum`` or ``mean``).
    """
 
    def __init__(
        self,
        alpha_mass=0,
        ignore_index=-1,
        reduction="mean",
    ):
        """
        Initialization
        """
        super().__init__()
 
        
        self.alpha_massalpha_mass = alpha_mass
 
        
        self.LCA_CELCA_CE = nn.CrossEntropyLoss(
            ignore_index=ignore_index, reduction=reduction
        )
        
        self.mass_CEmass_CE = nn.CrossEntropyLoss(
            ignore_index=ignore_index, reduction=reduction
        )
 
        assert alpha_mass >= 0, "Alpha should be positive"
 
    def forward(self, x_input, x_target, edge_input, edge_target, u_input, u_target):
        """
        Called internally by PyTorch to propagate the input.
        """
 
        LCA_loss = self.LCA_CELCA_CE(
            edge_input,
            edge_target,
        )
 
        mass_loss = (
            self.mass_CEmass_CE(
                x_input,
                x_target,
            )
            if self.alpha_massalpha_mass > 0
            else 0
        )
 
        return LCA_loss + self.alpha_massalpha_mass * mass_loss