PyTorch Lightning Result Example

pl_result_ex.py

    def training_step(self, batch: DataPoint, batch_idx: int) -> pl.TrainResult:
        planes, target_move, target_val = batch
        pred_move, pred_val = self(planes)
        mse_loss, cross_entropy_loss, loss = self.loss(
            pred_move, pred_val, target_move, target_val
        )
        result = pl.TrainResult(minimize=loss)
        result.log("train_loss", loss, prog_bar=True)
        result.log_dict(
            {
                "train_mse_loss": mse_loss,
                "train_ce_loss": cross_entropy_loss,
                "train_acc": accuracy(pred_move, target_move),
            }
        )
        return result

    def validation_step(self, batch: DataPoint, batch_idx: int) -> pl.EvalResult:
        planes, target_move, target_val = batch
        pred_move, pred_val = self(planes)
        mse_loss, cross_entropy_loss, loss = self.loss(
            pred_move, pred_val, target_move, target_val
        )
        result = pl.EvalResult(checkpoint_on=loss)
        result.log_dict(
            {
                "val_loss": loss,
                "val_mse_loss": mse_loss,
                "val_ce_loss": cross_entropy_loss,
                "val_acc": accuracy(pred_move, target_move),
            }
        )
        return result

    def test_step(self, batch: DataPoint, batch_idx: int) -> pl.EvalResult:
        result = self.validation_step(batch, batch_idx)
        result.rename_keys(
            {
                "val_loss": "test_loss",
                "val_mse_loss": "test_mse_loss",
                "val_ce_loss": "test_ce_loss",
                "val_acc": "test_acc",
            }
        )
        return result