| import math | |
| from typing import Dict, List, Optional | |
| import torch | |
| import torch.nn as nn | |
| class BundledLoss(nn.Module): | |
| def __init__( | |
| self, | |
| single_modality_loss, | |
| multi_view_consistency_loss, | |
| volume_mask_loss, | |
| multi_view_consistency_weight: float, | |
| mvc_time_dependent: bool, | |
| mvc_steepness: float, | |
| modality: List, | |
| consistency_weight: float, | |
| consistency_source: str, | |
| ): | |
| super().__init__() | |
| self.single_modality_loss = single_modality_loss | |
| self.multi_view_consistency_loss = multi_view_consistency_loss | |
| self.volume_mask_loss = volume_mask_loss | |
| self.mvc_weight = multi_view_consistency_weight | |
| self.mvc_time_dependent = mvc_time_dependent | |
| self.mvc_steepness = mvc_steepness | |
| self.modality = modality | |
| self.consistency_weight = consistency_weight | |
| self.consistency_source = consistency_source | |
| def forward( | |
| self, | |
| output: Dict, | |
| label, | |
| mask, | |
| epoch: int = 1, | |
| max_epoch: int = 70, | |
| spixel=None, | |
| raw_image=None, | |
| ): | |
| total_loss = 0.0 | |
| loss_dict = {} | |
| for modality in self.modality: | |
| single_loss = self.single_modality_loss(output[modality], label, mask) | |
| for k, v in single_loss.items(): | |
| loss_dict[f"{k}/{modality}"] = v | |
| total_loss = total_loss + single_loss["total_loss"] | |
| if self.mvc_time_dependent: | |
| mvc_weight = self.mvc_weight * math.exp( | |
| -self.mvc_steepness * (1 - epoch / max_epoch) ** 2 | |
| ) | |
| else: | |
| mvc_weight = self.mvc_weight | |
| multi_view_consistency_loss = self.multi_view_consistency_loss( | |
| output, label, spixel, raw_image, mask | |
| ) | |
| for k, v in multi_view_consistency_loss.items(): | |
| if k not in ["total_loss", "tgt_map"]: | |
| loss_dict.update({k: v}) | |
| if self.consistency_weight != 0.0 and self.consistency_source == "ensemble": | |
| for modality in self.modality: | |
| consisitency_loss = self.volume_mask_loss( | |
| output[modality]["out_vol"], multi_view_consistency_loss["tgt_map"] | |
| ) | |
| consisitency_loss = consisitency_loss["loss"] | |
| loss_dict[f"consistency_loss/{modality}"] = consisitency_loss | |
| total_loss = ( | |
| total_loss | |
| + self.consistency_weight | |
| * consisitency_loss | |
| * math.exp(-self.mvc_steepness * (1 - epoch / max_epoch) ** 2) | |
| ) | |
| total_loss = total_loss + mvc_weight * multi_view_consistency_loss["total_loss"] | |
| return {"total_loss": total_loss, **loss_dict} | |