BasePredictionWriter

class pytorch_lightning.callbacks.BasePredictionWriter(write_interval='batch')

Bases: pytorch_lightning.callbacks.callback.Callback

Base class to implement how the predictions should be stored.

Parameters:

write_interval (str) – When to write.

Example:

import torch
from pytorch_lightning.callbacks import BasePredictionWriter

class CustomWriter(BasePredictionWriter):

    def __init__(self, output_dir: str, write_interval: str):
        super().__init__(write_interval)
        self.output_dir = output_dir

    def write_on_batch_end(
        self, trainer, pl_module: 'LightningModule', prediction: Any, batch_indices: List[int], batch: Any,
        batch_idx: int, dataloader_idx: int
    ):
        torch.save(prediction, os.path.join(self.output_dir, dataloader_idx, f"{batch_idx}.pt"))

    def write_on_epoch_end(
        self, trainer, pl_module: 'LightningModule', predictions: List[Any], batch_indices: List[Any]
    ):
        torch.save(predictions, os.path.join(self.output_dir, "predictions.pt"))

on_predict_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)

Called when the predict batch ends.

Return type:

None

on_predict_epoch_end(trainer, pl_module, outputs)

Called when the predict epoch ends.

Return type:

None

write_on_batch_end(trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx)

Override with the logic to write a single batch.

Return type:

None

write_on_epoch_end(trainer, pl_module, predictions, batch_indices)

Override with the logic to write all batches.

Return type:

None