mlflow.pytorch

The mlflow.pytorch module provides an API for logging and loading PyTorch models. This module exports PyTorch models with the following flavors:

PyTorch (native) format
This is the main flavor that can be loaded back into PyTorch.
mlflow.pyfunc
Produced for use by generic pyfunc-based deployment tools and batch inference.
mlflow.pytorch.load_model(path, run_id=None, **kwargs)

Load a PyTorch model from a local file (if run_id is None) or a run.

Parameters:
  • path – Local filesystem path or run-relative artifact path to the model saved by mlflow.pytorch.log_model().
  • run_id – Run ID. If provided, combined with path to identify the model.
  • kwargs – kwargs to pass to torch.load method.
>>> import torch
>>> import mlflow
>>> import mlflow.pytorch
>>> # set values
>>> model_path_dir = ...
>>> run_id="96771d893a5e46159d9f3b49bf9013e2"
>>> pytorch_model = mlflow.pytorch.load_model(model_path_dir, run_id)
>>> y_pred = pytorch_model(x_new_data)
mlflow.pytorch.log_model(pytorch_model, artifact_path, conda_env=None, code_paths=None, **kwargs)

Log a PyTorch model as an MLflow artifact for the current run.

Parameters:
  • pytorch_model

    PyTorch model to be saved. Must accept a single torch.FloatTensor as input and produce a single output tensor. Any code dependencies of the model’s class, including the class definition itself, should be included in one of the following locations:

    • The package(s) listed in the model’s Conda environment, specified by the conda_env parameter.
    • One or more of the files specified by the code_paths parameter.
  • artifact_path – Run-relative artifact path.
  • conda_env

    Path to a Conda environment file. If provided, this decribes the environment this model should be run in. At minimum, it should specify the dependencies contained in mlflow.pytorch.DEFAULT_CONDA_ENV. If None, the default mlflow.pytorch.DEFAULT_CONDA_ENV environment will be added to the model. The following is an example dictionary representation of a Conda environment:

    {
        'name': 'mlflow-env',
        'channels': ['defaults'],
        'dependencies': [
            'python=3.7.0',
            'pytorch=0.4.1',
            'torchvision=0.2.1'
        ]
    }
    
  • code_paths – A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files will be prepended to the system path when the model is loaded.
  • kwargs – kwargs to pass to torch.save method.
>>> import torch
>>> import mlflow
>>> import mlflow.pytorch
>>> # X data
>>> x_data = torch.Tensor([[1.0], [2.0], [3.0]])
>>> # Y data with its expected value: labels
>>> y_data = torch.Tensor([[2.0], [4.0], [6.0]])
>>> # Partial Model example modified from Sung Kim
>>> # https://github.com/hunkim/PyTorchZeroToAll
>>> class Model(torch.nn.Module):
>>>    def __init__(self):
>>>       super(Model, self).__init__()
>>>       self.linear = torch.nn.Linear(1, 1)  # One in and one out
>>>    def forward(self, x):
>>>        y_pred = self.linear(x)
>>>        return y_pred
>>> # our model
>>> model = Model()
>>> criterion = torch.nn.MSELoss(size_average=False)
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
>>> # Training loop
>>> for epoch in range(500):
>>>    # Forward pass: Compute predicted y by passing x to the model
>>>    y_pred = model(x_data)
>>>    # Compute and print loss
>>>    loss = criterion(y_pred, y_data)
>>>    print(epoch, loss.data[0])
>>>   #Zero gradients, perform a backward pass, and update the weights.
>>>   optimizer.zero_grad()
>>>   loss.backward()
>>>   optimizer.step()
>>> # After training
>>> for hv in [4.0, 5.0, 6.0]:
>>>     hour_var = torch.Tensor([[hv]])
>>>     y_pred = model(hour_var)
>>>     print("predict (after training)",  hv, model(hour_var).data[0][0])
>>> # log the model
>>> with mlflow.start_run() as run:
>>>   mlflow.log_param("epochs", 500)
>>>   mlflow.pytorch.log_model(pytorch_model, "models")
mlflow.pytorch.save_model(pytorch_model, path, conda_env=None, mlflow_model=<mlflow.models.Model object>, code_paths=None, **kwargs)

Save a PyTorch model to a path on the local file system.

Parameters:
  • pytorch_model

    PyTorch model to be saved. Must accept a single torch.FloatTensor as input and produce a single output tensor. Any code dependencies of the model’s class, including the class definition itself, should be included in one of the following locations:

    • The package(s) listed in the model’s Conda environment, specified by the conda_env parameter.
    • One or more of the files specified by the code_paths parameter.
  • path – Local path where the model is to be saved.
  • conda_env

    Either a dictionary representation of a Conda environment or the path to a Conda environment yaml file. If provided, this decribes the environment this model should be run in. At minimum, it should specify the dependencies contained in mlflow.pytorch.DEFAULT_CONDA_ENV. If None, the default mlflow.pytorch.DEFAULT_CONDA_ENV environment will be added to the model. The following is an example dictionary representation of a Conda environment:

    {
        'name': 'mlflow-env',
        'channels': ['defaults'],
        'dependencies': [
            'python=3.7.0',
            'pytorch=0.4.1',
            'torchvision=0.2.1'
        ]
    }
    
  • mlflow_modelmlflow.models.Model this flavor is being added to.
  • code_paths – A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files will be prepended to the system path when the model is loaded.
  • kwargs – kwargs to pass to torch.save method.
>>> import torch
>>> import mlflow
>>> import mlflow.pytorch
>>> # create model and set values
>>> pytorch_model = Model()
>>> pytorch_model_path = ...
>>> #train our model
>>> for epoch in range(500):
>>>     y_pred = model(x_data)
>>>     ...
>>> #save the model
>>> with mlflow.start_run() as run:
>>>   mlflow.log_param("epochs", 500)
>>>   mlflow.pytorch.save_model(pytorch_model, pytorch_model_path)