PyTorch Lightning#
Here’s a simple example of using PyTorch Lightning with BentoML:
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_model("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_model("addition_model:latest")
# Run a given model under `Runner` abstraction with `to_runner`
runner = bentoml.pytorch_lightning.get(tag).to_runner()
runner.init_local()
runner.run(torch.from_numpy(np.array([[1,2,3,4]])))
Note
You can find more examples for PyTorch Lightning in our bentoml/examples directory.
- bentoml.pytorch_lightning.save_model(name: Tag | str, model: pl.LightningModule, *, signatures: ModelSignaturesType | None = None, labels: t.Dict[str, str] | None = None, custom_objects: t.Dict[str, t.Any] | None = None, external_modules: t.List[ModuleType] | None = None, metadata: t.Dict[str, t.Any] | None = None) bentoml.Model [source]#
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 jsonexternal_modules (
List[ModuleType]
, optional, default toNone
) – user-defined additional python modules to be saved alongside the model or custom objects, e.g. a tokenizer module, preprocessor module, model configuration modulemetadata (
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_model("lit_classifier", LitClassifier())
- bentoml.pytorch_lightning.load_model(bentoml_model: str | Tag | Model, device_id: str | None = 'cpu') torch.ScriptModule [source]#
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. Refer 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
torch.ScriptModule
from BentoML modelstore.- Return type:
torch.ScriptModule
Examples:
import bentoml lit = bentoml.torchscript.load_model('lit_classifier:latest', device_id="cuda:0")