| import typing as tp | |
| from torch.nn import functional as F | |
| from torch import nn | |
| import torch | |
| class LossModule(nn.Module): | |
| def __init__(self, name: str, weight: float = 1.0): | |
| super().__init__() | |
| self.name = name | |
| self.weight = weight | |
| def forward(self, info, *args, **kwargs): | |
| raise NotImplementedError | |
| class ValueLoss(LossModule): | |
| def __init__(self, key: str, name, weight: float = 1.0): | |
| super().__init__(name=name, weight=weight) | |
| self.key = key | |
| def forward(self, info): | |
| return self.weight * info[self.key] | |
| class L1Loss(LossModule): | |
| def __init__(self, key_a: str, key_b: str, weight: float = 1.0, mask_key: str = None, name: str = 'l1_loss'): | |
| super().__init__(name=name, weight=weight) | |
| self.key_a = key_a | |
| self.key_b = key_b | |
| self.mask_key = mask_key | |
| def forward(self, info): | |
| mse_loss = F.l1_loss(info[self.key_a], info[self.key_b], reduction='none') | |
| if self.mask_key is not None and self.mask_key in info: | |
| mse_loss = mse_loss[info[self.mask_key]] | |
| mse_loss = mse_loss.mean() | |
| return self.weight * mse_loss | |
| class MSELoss(LossModule): | |
| def __init__(self, key_a: str, key_b: str, weight: float = 1.0, mask_key: str = None, name: str = 'mse_loss'): | |
| super().__init__(name=name, weight=weight) | |
| self.key_a = key_a | |
| self.key_b = key_b | |
| self.mask_key = mask_key | |
| def forward(self, info): | |
| mse_loss = F.mse_loss(info[self.key_a], info[self.key_b], reduction='none') | |
| if self.mask_key is not None and self.mask_key in info and info[self.mask_key] is not None: | |
| mask = info[self.mask_key] | |
| if mask.ndim == 2 and mse_loss.ndim == 3: | |
| mask = mask.unsqueeze(1) | |
| if mask.shape[1] != mse_loss.shape[1]: | |
| mask = mask.repeat(1, mse_loss.shape[1], 1) | |
| mse_loss = mse_loss[mask] | |
| mse_loss = mse_loss.mean() | |
| return self.weight * mse_loss | |
| class AuralossLoss(LossModule): | |
| def __init__(self, auraloss_module, input_key: str, target_key: str, name: str, weight: float = 1): | |
| super().__init__(name, weight) | |
| self.auraloss_module = auraloss_module | |
| self.input_key = input_key | |
| self.target_key = target_key | |
| def forward(self, info): | |
| loss = self.auraloss_module(info[self.input_key], info[self.target_key]) | |
| return self.weight * loss | |
| class MultiLoss(nn.Module): | |
| def __init__(self, losses: tp.List[LossModule]): | |
| super().__init__() | |
| self.losses = nn.ModuleList(losses) | |
| def forward(self, info): | |
| total_loss = 0 | |
| losses = {} | |
| for loss_module in self.losses: | |
| module_loss = loss_module(info) | |
| total_loss += module_loss | |
| losses[loss_module.name] = module_loss | |
| return total_loss, losses |