Storage Format
There are two types of The MLflow Models.
- serialized object (e.g., a picled
scikit-learnmodel) - Python Script (or notebook)
Signature
from mlflow.models import infer_signature train = df.drop_column("target_label") predictions = ... # compute model predictions signature = infer_signature(train, predictions) log_model vs save_model
| Function | Purpose |
|---|---|
log_model | Stores model in MLflow tracking server with experiment metadata |
save_model | Saves model to local filesystem only without tracking |
log_modelintegrates with experiment tracking,save_modelis standalone storage.
Onnx example
import numpy as np import mlflow from mlflow.models import infer_signature import onnx import torch from torch import nn # define a torch model net = nn.Linear(6, 1) loss_function = nn.L1Loss() optimizer = torch.optim.Adam(net.parameters(), lr=1e-4) X = torch.randn(6) y = torch.randn(1) # run model training epochs = 5 for epoch in range(epochs): optimizer.zero_grad() outputs = net(X) loss = loss_function(outputs, y) loss.backward() optimizer.step() # convert model to ONNX and load it torch.onnx.export(net, X, "model.onnx") onnx_model = onnx.load_model("model.onnx") # log the model into a mlflow run with mlflow.start_run(): signature = infer_signature(X.numpy(), net(X).detach().numpy()) model_info = mlflow.onnx.log_model(onnx_model, "model", signature=signature) # load the logged model and make a prediction onnx_pyfunc = mlflow.pyfunc.load_model(model_info.model_uri) predictions = onnx_pyfunc.predict(X.numpy()) print(predictions) Pytorch Example
import numpy as np import mlflow from mlflow.models import infer_signature import torch from torch import nn net = nn.Linear(6, 1) loss_function = nn.L1Loss() optimizer = torch.optim.Adam(net.parameters(), lr=1e-4) X = torch.randn(6) y = torch.randn(1) epochs = 5 for epoch in range(epochs): optimizer.zero_grad() outputs = net(X) loss = loss_function(outputs, y) loss.backward() optimizer.step() with mlflow.start_run() as run: signature = infer_signature(X.numpy(), net(X).detach().numpy()) model_info = mlflow.pytorch.log_model(net, "model", signature=signature) pytorch_pyfunc = mlflow.pyfunc.load_model(model_uri=model_info.model_uri) predictions = pytorch_pyfunc.predict(torch.randn(6).numpy()) print(predictions) 
Top comments (0)