PyTorch Lightning

Users can now use PyTorch Lightning with BentoML with the following API: load, save, and load_runner as follow:

import bentoml
import torch
import pytorch_lightning as pl

class AdditionModel(pl.LightningModule):
   def forward(self, inputs):
      return inputs.add(1)

# `save` a given classifier and retrieve coresponding tag:
tag = bentoml.pytorch_lightning.save("addition_model", AdditionModel())

# retrieve metadata with `bentoml.models.get`:
metadata = bentoml.models.get(tag)

# `load` the model back in memory:
model = bentoml.pytorch_lightning.load("addition_model:latest")

# Run a given model under `Runner` abstraction with `load_runner`
runner = bentoml.pytorch_lightning.load_runner(tag)
runner.run_batch(torch.from_numpy(np.array([[1,2,3,4]])))

Note

You can find more examples for PyTorch Lightning in our gallery repo.

bentoml.pytorch_lightning.save(name, model, *, labels=None, custom_objects=None, metadata=None)

Save a model instance to BentoML modelstore.

Parameters

  • name (str) – Name for given model instance. This should pass Python identifier check.

  • model (pl.LightningModule) – Instance of model to be saved

  • labels (Dict[str, str], optional, default to None) – user-defined labels for managing models, e.g. team=nlp, stage=dev

  • custom_objects (Dict[str, Any]], optional, default to None) – user-defined additional python objects to be saved alongside the model, e.g. a tokenizer instance, preprocessor function, model configuration json

  • metadata (Dict[str, Any], optional, default to None) – Custom metadata for given model.

  • model_store (ModelStore, default to BentoMLContainer.model_store) – BentoML modelstore, provided by DI Container.

Returns

A tag with a format name:version where name is the user-defined model’s name, and a generated version by BentoML.

Return type

Tag

Examples:

import bentoml
import torch
import pytorch_lightning as pl

class LitClassifier(pl.LightningModule):

    def __init__(self, hidden_dim: int = 128, learning_rate: float = 0.0001):
        super().__init__()
        self.save_hyperparameters()

        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

tag = bentoml.pytorch_lightning.save("lit_classifier", LitClassifier())
bentoml.pytorch_lightning.load(tag, device_id='cpu', model_store=<simple_di.providers.SingletonFactory object>)

Load a model from BentoML local modelstore with given name.

Parameters
Returns

an instance of pl.LightningModule from BentoML modelstore.

Return type

pl.LightningModule

Examples:

import bentoml
lit = bentoml.pytorch_lightning.load('lit_classifier:latest', device_id="cuda:0")
bentoml.pytorch_lightning.load_runner(tag, *, predict_fn_name='__call__', partial_kwargs=None, name=None)

Runner represents a unit of serving logic that can be scaled horizontally to maximize throughput. bentoml.pytorch_lightning.load_runner() implements a Runner class that wrap around a pl.LightningModule instance, which optimize it for the BentoML runtime.

Parameters

  • tag (Union[str, Tag]) – Tag of a saved model in BentoML local modelstore.

  • predict_fn_name (str, default to __call__) – inference function to be used.

  • partial_kwargs (Dict[str, Any], optional, default to None) – Common kwargs passed to model for this runner

Returns

Runner instances for bentoml.pytorch_lightning model

Return type

Runner

Examples:

import bentoml.pytorch_lightning
runner = bentoml.pytorch_lightning.load_runner("lit_classifier:20201012_DE43A2")
runner.run(pd.DataFrame("/path/to/csv"))