class ray.train.lightning.LightningPredictor(model: pytorch_lightning.core.lightning.LightningModule, preprocessor: Optional[ray.data.preprocessor.Preprocessor] = None, use_gpu: bool = False)[source]#

Bases: ray.train.torch.torch_predictor.TorchPredictor

A predictor for PyTorch Lightning modules.


import torch
import numpy as np
import pytorch_lightning as pl
from ray.train.lightning import LightningPredictor

class MyModel(pl.LightningModule):
    def __init__(self, input_dim, output_dim):
        self.linear = torch.nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = self.linear(x)
        return out

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = torch.nn.functional.mse_loss(y_hat, y)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

batch_size, input_dim, output_dim = 10, 3, 5
model = MyModel(input_dim=input_dim, output_dim=output_dim)
predictor = LightningPredictor(model=model, use_gpu=False)
batch = np.random.rand(batch_size, input_dim).astype(np.float32)

# Internally, LightningPredictor.predict()  invokes the forward() method
# of the model to generate predictions
output = predictor.predict(batch)

assert output["predictions"].shape == (batch_size, output_dim)
  • model – The PyTorch Lightning module to use for predictions.

  • preprocessor – A preprocessor used to transform data batches prior to prediction.

  • use_gpu – If set, the model will be moved to GPU on instantiation and prediction happens on GPU.

PublicAPI (alpha): This API is in alpha and may change before becoming stable.



Runs inference on a single batch of tensor data.

from_checkpoint(checkpoint, model_class, *)

Instantiate the LightningPredictor from a Checkpoint.


Create a Predictor from a Pandas UDF.


Get the preprocessor to use prior to executing predictions.

predict(data[, dtype])

Run inference on data batch.


DeveloperAPI: This API may change across minor Ray releases.


Set the preprocessor to use prior to executing predictions.