Torchmetric that grabs the precomputed ce_loss value from the model outputs
| 81 | |
| 82 | @rename_class("LanguageCrossEntropy") |
| 83 | class EfficientCrossEntropy(Metric): |
| 84 | """Torchmetric that grabs the precomputed ce_loss value from the model outputs""" |
| 85 | |
| 86 | # Make torchmetrics call update only once |
| 87 | full_state_update = False |
| 88 | |
| 89 | def __init__(self, dist_sync_on_step: bool = False): |
| 90 | super().__init__(dist_sync_on_step=dist_sync_on_step) |
| 91 | self.add_state("sum_loss", default=torch.tensor(0.0), dist_reduce_fx="sum") |
| 92 | self.add_state("total_items", default=torch.tensor(0), dist_reduce_fx="sum") |
| 93 | |
| 94 | def update(self, loss: Tensor) -> None: |
| 95 | """Updates the internal state with results from a new batch. |
| 96 | |
| 97 | Args: |
| 98 | loss (~torch.Tensor): A Tensor of loss values to compare against. |
| 99 | """ |
| 100 | self.sum_loss += loss |
| 101 | self.total_items += 1 |
| 102 | |
| 103 | def compute(self) -> Tensor: |
| 104 | """Aggregate the state over all processes to compute the metric. |
| 105 | |
| 106 | Returns: |
| 107 | loss: The loss averaged across all batches as a :class:`~torch.Tensor`. |
| 108 | """ |
| 109 | # Return average loss over entire dataset |
| 110 | return self.sum_loss / self.total_items # type: ignore (third-party) |
| 111 | |
| 112 | |
| 113 | @rename_class("ZLoss") |
no outgoing calls
no test coverage detected