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 toNone
) – user-defined labels for managing models, e.g. team=nlp, stage=devcustom_objects (
Dict[str, Any]]
, optional, default toNone
) – user-defined additional python objects to be saved alongside the model, e.g. a tokenizer instance, preprocessor function, model configuration jsonmetadata (
Dict[str, Any]
, optional, default toNone
) – Custom metadata for given model.model_store (
ModelStore
, default toBentoMLContainer.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
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
tag (
Union[str, Tag]
) – Tag of a saved model in BentoML local modelstore.device_id (
str
, optional) – Optional devices to put the given model on. Refers to https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.devicemodel_store (
ModelStore
, default toBentoMLContainer.model_store
) – BentoML modelstore, provided by DI Container.
- 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 apl.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 toNone
) – 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"))