(self, batch: BatchedLabeledPathContext, step: str)
| 53 | # ========== MODEL STEP ========== |
| 54 | |
| 55 | def _shared_step(self, batch: BatchedLabeledPathContext, step: str) -> Dict: |
| 56 | # [batch size; num_classes] |
| 57 | logits = self(batch.from_token, batch.path_nodes, batch.to_token, batch.contexts_per_label) |
| 58 | labels = batch.labels.squeeze(0) |
| 59 | loss = torch.nn.functional.cross_entropy(logits, labels) |
| 60 | |
| 61 | with torch.no_grad(): |
| 62 | predictions = logits.argmax(-1) |
| 63 | accuracy = self.__metrics[f"{step}_acc"](predictions, labels) |
| 64 | |
| 65 | return {f"{step}/loss": loss, f"{step}/accuracy": accuracy} |
| 66 | |
| 67 | def training_step(self, batch: BatchedLabeledPathContext, batch_idx: int) -> Dict: # type: ignore |
| 68 | result = self._shared_step(batch, "train") |
no outgoing calls