def load_model(model_uri, dst_path=None):
"""
Load an ONNX model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model, for example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see the
`Artifacts Documentation <https://www.mlflow.org/docs/latest/
tracking.html#artifact-stores>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: An ONNX model instance.
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
onnx_model_artifacts_path = os.path.join(local_model_path,
flavor_conf["data"])
return _load_model(model_file=onnx_model_artifacts_path)
def load_model(model_uri, model=None, dst_path=None, **kwargs):
"""
Load a paddle model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model, for example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
:param model: Required when loading a `paddle.Model` model saved with `training=True`.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:param kwargs: The keyword arguments to pass to `paddle.jit.load`
or `model.load`.
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:return: A paddle model.
.. code-block:: python
:caption: Example
import mlflow.paddle
pd_model = mlflow.paddle.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2/pd_models")
# use Pandas DataFrame to make predictions
np_array = ...
predictions = pd_model(np_array)
"""
import paddle
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
pd_model_artifacts_path = os.path.join(local_model_path,
flavor_conf["pickled_model"])
if model is None:
return paddle.jit.load(pd_model_artifacts_path, **kwargs)
elif not isinstance(model, paddle.Model):
raise TypeError(
"Invalid object type `{}` for `model`, must be `paddle.Model`".
format(type(model)))
else:
contains_pdparams = _contains_pdparams(local_model_path)
if not contains_pdparams:
raise TypeError(
"This model can't be loaded via `model.load` because a '.pdparams' file "
"doesn't exist. Please leave `model` unspecified to load the model via "
"`paddle.jit.load` or set `training` to True when saving a model."
)
model.load(pd_model_artifacts_path, **kwargs)
return model
def load_model(model_uri, dst_path=None):
"""
Load a CatBoost model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: A CatBoost model (an instance of `CatBoost`_, `CatBoostClassifier`_,
or `CatBoostRegressor`_)
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
cb_model_file_path = os.path.join(
local_model_path,
flavor_conf.get(_MODEL_BINARY_KEY, _MODEL_BINARY_FILE_NAME))
return _load_model(cb_model_file_path, flavor_conf.get(_MODEL_TYPE_KEY),
flavor_conf.get(_SAVE_FORMAT_KEY))
def load_model(model_uri, dst_path=None):
"""
Load an XGBoost model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: An XGBoost model. An instance of either `xgboost.Booster`_ or XGBoost scikit-learn
models, depending on the saved model class specification.
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(local_model_path, FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
return _load_model(path=local_model_path)
def load_model(model_uri, dst_path=None):
"""
Load a spaCy model from a local file (if ``run_id`` is ``None``) or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: A spaCy loaded model
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
# Flavor configurations for models saved in MLflow version <= 0.8.0 may not contain a
# `data` key; in this case, we assume the model artifact path to be `model.spacy`
spacy_model_file_path = os.path.join(
local_model_path, flavor_conf.get("data", "model.spacy"))
return _load_model(path=spacy_model_file_path)
def load_model(model_uri, dst_path=None):
"""
Load a statsmodels model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: A statsmodels model (an instance of `statsmodels.base.model.Results`_).
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
statsmodels_model_file_path = os.path.join(
local_model_path, flavor_conf.get("data", STATSMODELS_DATA_SUBPATH))
return _load_model(path=statsmodels_model_file_path)
def load_explainer(model_uri):
"""
Load a SHAP explainer from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model, for example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:return: A SHAP explainer.
"""
explainer_path = _download_artifact_from_uri(artifact_uri=model_uri)
flavor_conf = _get_flavor_configuration(model_path=explainer_path, flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(explainer_path, flavor_conf)
explainer_artifacts_path = os.path.join(explainer_path, flavor_conf["serialized_explainer"])
underlying_model_flavor = flavor_conf["underlying_model_flavor"]
model = None
if underlying_model_flavor != _UNKNOWN_MODEL_FLAVOR:
underlying_model_path = os.path.join(explainer_path, _UNDERLYING_MODEL_SUBPATH)
if underlying_model_flavor == mlflow.sklearn.FLAVOR_NAME:
model = mlflow.sklearn._load_pyfunc(underlying_model_path).predict
elif underlying_model_flavor == mlflow.pytorch.FLAVOR_NAME:
model = mlflow.pytorch._load_model(os.path.join(underlying_model_path, "data"))
return _load_explainer(explainer_file=explainer_artifacts_path, model=model)
def load_model(model_uri, ctx, dst_path=None):
"""
Load a Gluon model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:param ctx: Either CPU or GPU.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: A Gluon model instance.
.. code-block:: python
:caption: Example
# Load persisted model as a Gluon model, make inferences against an NDArray
model = mlflow.gluon.load_model("runs:/" + gluon_random_data_run.info.run_id + "/model")
model(nd.array(np.random.rand(1000, 1, 32)))
"""
import mxnet as mx
from mxnet import gluon
from mxnet import sym
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
model_arch_path = os.path.join(local_model_path, "data",
_MODEL_SAVE_PATH) + "-symbol.json"
model_params_path = os.path.join(local_model_path, "data",
_MODEL_SAVE_PATH) + "-0000.params"
if Version(mx.__version__) >= Version("2.0.0"):
return gluon.SymbolBlock.imports(model_arch_path,
input_names=["data"],
param_file=model_params_path,
ctx=ctx)
else:
symbol = sym.load(model_arch_path)
inputs = sym.var("data", dtype="float32")
net = gluon.SymbolBlock(symbol, inputs)
net.collect_params().load(model_params_path, ctx)
return net
def test_add_code_to_system_path(sklearn_knn_model, model_path):
mlflow.sklearn.save_model(
sk_model=sklearn_knn_model,
path=model_path,
code_paths=["tests/utils/test_resources/dummy_module.py"],
)
sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME)
with pytest.raises(ModuleNotFoundError,
match="No module named 'dummy_module'"):
import dummy_module
mlflow_model_utils._add_code_from_conf_to_system_path(
model_path, sklearn_flavor_config)
import dummy_module
def load_model(model_uri, dfs_tmpdir=None):
"""
Load the Spark MLlib model from the path.
:param model_uri: The location, in URI format, of the MLflow model, for example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:param dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
filesystem if running in local mode. The model is loaded from this
destination. Defaults to ``/tmp/mlflow``.
:return: pyspark.ml.pipeline.PipelineModel
.. code-block:: python
:caption: Example
from mlflow import spark
model = mlflow.spark.load_model("spark-model")
# Prepare test documents, which are unlabeled (id, text) tuples.
test = spark.createDataFrame([
(4, "spark i j k"),
(5, "l m n"),
(6, "spark hadoop spark"),
(7, "apache hadoop")], ["id", "text"])
# Make predictions on test documents
prediction = model.transform(test)
"""
if RunsArtifactRepository.is_runs_uri(model_uri):
runs_uri = model_uri
model_uri = RunsArtifactRepository.get_underlying_uri(model_uri)
_logger.info("'%s' resolved as '%s'", runs_uri, model_uri)
elif ModelsArtifactRepository.is_models_uri(model_uri):
runs_uri = model_uri
model_uri = ModelsArtifactRepository.get_underlying_uri(model_uri)
_logger.info("'%s' resolved as '%s'", runs_uri, model_uri)
flavor_conf = _get_flavor_configuration_from_uri(model_uri, FLAVOR_NAME)
model_uri = append_to_uri_path(model_uri, flavor_conf["model_data"])
local_model_path = _download_artifact_from_uri(model_uri)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
return _load_model(model_uri=model_uri, dfs_tmpdir_base=dfs_tmpdir)
def load_model(model_uri, dst_path=None, **kwargs):
"""
Load a Keras model from a local file or a run.
Extra arguments are passed through to keras.load_model.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: A Keras model instance.
.. code-block:: python
:caption: Example
# Load persisted model as a Keras model or as a PyFunc, call predict() on a pandas DataFrame
keras_model = mlflow.keras.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2" + "/models")
predictions = keras_model.predict(x_test)
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
keras_module = importlib.import_module(
flavor_conf.get("keras_module", "keras"))
keras_model_artifacts_path = os.path.join(
local_model_path, flavor_conf.get("data", _MODEL_SAVE_PATH))
# For backwards compatibility, we assume h5 when the save_format is absent
save_format = flavor_conf.get("save_format", "h5")
return _load_model(
model_path=keras_model_artifacts_path,
keras_module=keras_module,
save_format=save_format,
**kwargs,
)
def _load_model(path):
import diviner
local_path = pathlib.Path(path)
flavor_conf = _get_flavor_configuration(model_path=str(local_path),
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(str(local_path), flavor_conf)
diviner_model_path = local_path.joinpath(
flavor_conf.get(_MODEL_BINARY_KEY, _MODEL_BINARY_FILE_NAME))
diviner_instance = getattr(diviner, _get_diviner_instance_type(path))
return diviner_instance.load(str(diviner_model_path))
def test_add_code_to_system_path(sklearn_knn_model, model_path):
mlflow.sklearn.save_model(
sk_model=sklearn_knn_model,
path=model_path,
code_paths=["tests/utils/test_resources/dummy_module.py"],
)
sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
)
with TempDir(chdr=True):
# Load the model from a new directory that is not a parent of the source code path to
# verify that source code paths and their subdirectories are resolved correctly
with pytest.raises(ModuleNotFoundError, match="No module named 'dummy_module'"):
import dummy_module
mlflow_model_utils._add_code_from_conf_to_system_path(model_path, sklearn_flavor_config)
import dummy_module
def load_model(model_uri, dst_path=None):
"""
Load a fastai model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:return: A fastai model (an instance of `fastai.Learner`_).
.. code-block:: python
:caption: Example
import mlflow.fastai
# Define the Learner model
model = ...
# log the fastai Leaner model
with mlflow.start_run() as run:
model.fit(epochs, learning_rate)
mlflow.fastai.log_model(model, "model")
# Load the model for scoring
model_uri = "runs:/{}/model".format(run.info.run_id)
loaded_model = mlflow.fastai.load_model(model_uri)
results = loaded_model.predict(predict_data)
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, flavor_conf)
model_file_path = os.path.join(local_model_path,
flavor_conf.get("data", "model.fastai"))
return _load_model(path=model_file_path)
def load_model(model_uri, dst_path=None, **kwargs):
"""
Load a PyTorch model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model, for example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:param dst_path: The local filesystem path to which to download the model artifact.
This directory must already exist. If unspecified, a local output
path will be created.
:param kwargs: kwargs to pass to ``torch.load`` method.
:return: A PyTorch model.
.. code-block:: python
:caption: Example
import torch
import mlflow.pytorch
# Class defined here
class LinearNNModel(torch.nn.Module):
...
# Initialize our model, criterion and optimizer
...
# Training loop
...
# Log the model
with mlflow.start_run() as run:
mlflow.pytorch.log_model(model, "model")
# Inference after loading the logged model
model_uri = "runs:/{}/model".format(run.info.run_id)
loaded_model = mlflow.pytorch.load_model(model_uri)
for x in [4.0, 6.0, 30.0]:
X = torch.Tensor([[x]])
y_pred = loaded_model(X)
print("predict X: {}, y_pred: {:.2f}".format(x, y_pred.data.item()))
.. code-block:: text
:caption: Output
predict X: 4.0, y_pred: 7.57
predict X: 6.0, y_pred: 11.64
predict X: 30.0, y_pred: 60.48
"""
import torch
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri,
output_path=dst_path)
pytorch_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(local_model_path, pytorch_conf)
if torch.__version__ != pytorch_conf["pytorch_version"]:
_logger.warning(
"Stored model version '%s' does not match installed PyTorch version '%s'",
pytorch_conf["pytorch_version"],
torch.__version__,
)
torch_model_artifacts_path = os.path.join(local_model_path,
pytorch_conf["model_data"])
return _load_model(path=torch_model_artifacts_path, **kwargs)
