| 16 | |
| 17 | |
| 18 | class Code2Class(LightningModule): |
| 19 | def __init__(self, model_config: DictConfig, optimizer_config: DictConfig, vocabulary: Vocabulary): |
| 20 | super().__init__() |
| 21 | self.save_hyperparameters() |
| 22 | self._optim_config = optimizer_config |
| 23 | |
| 24 | self._encoder = PathEncoder( |
| 25 | model_config, |
| 26 | len(vocabulary.token_to_id), |
| 27 | vocabulary.token_to_id[Vocabulary.PAD], |
| 28 | len(vocabulary.node_to_id), |
| 29 | vocabulary.node_to_id[Vocabulary.PAD], |
| 30 | ) |
| 31 | |
| 32 | self._classifier = Classifier(model_config, len(vocabulary.label_to_id)) |
| 33 | |
| 34 | metrics: Dict[str, Metric] = { |
| 35 | f"{holdout}_acc": Accuracy(num_classes=len(vocabulary.label_to_id)) for holdout in ["train", "val", "test"] |
| 36 | } |
| 37 | self.__metrics = MetricCollection(metrics) |
| 38 | |
| 39 | def configure_optimizers(self) -> Tuple[List[Optimizer], List[_LRScheduler]]: |
| 40 | return configure_optimizers_alon(self._optim_config, self.parameters()) |
| 41 | |
| 42 | def forward( # type: ignore |
| 43 | self, |
| 44 | from_token: torch.Tensor, |
| 45 | path_nodes: torch.Tensor, |
| 46 | to_token: torch.Tensor, |
| 47 | contexts_per_label: torch.Tensor, |
| 48 | ) -> torch.Tensor: |
| 49 | encoded_paths = self._encoder(from_token, path_nodes, to_token) |
| 50 | output_logits = self._classifier(encoded_paths, contexts_per_label) |
| 51 | return output_logits |
| 52 | |
| 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") |
| 69 | self.log_dict(result, on_step=True, on_epoch=False) |
| 70 | self.log("acc", result["train/accuracy"], prog_bar=True, logger=False) |
| 71 | return result["train/loss"] |
| 72 | |
| 73 | def validation_step(self, batch: BatchedLabeledPathContext, batch_idx: int) -> Dict: # type: ignore |
| 74 | return self._shared_step(batch, "val") |
| 75 | |