def test_missing_value_hint_is_displayed_when_it_should():
m = Model()
input_schema = Schema([ColSpec("integer", "a")])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
pdf = pd.DataFrame(
data=[[1], [None]],
columns=["a"],
)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
hint = "Hint: the type mismatch is likely caused by missing values."
assert "Incompatible input types" in str(ex.value.message)
assert hint in str(ex.value.message)
pdf = pd.DataFrame(
data=[[1.5], [None]],
columns=["a"],
)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
assert hint not in str(ex.value.message)
pdf = pd.DataFrame(data=[[1], [2]], columns=["a"], dtype=np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex.value.message)
assert hint not in str(ex.value.message)
def test_schema_enforcement_named_tensor_schema_1d():
m = Model()
input_schema = Schema([
TensorSpec(np.dtype(np.uint64), (-1, ), "a"),
TensorSpec(np.dtype(np.float32), (-1, ), "b")
])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
pdf = pd.DataFrame(data=[[0, 0], [1, 1]], columns=["a", "b"])
pdf["a"] = pdf["a"].astype(np.uint64)
pdf["b"] = pdf["a"].astype(np.float32)
d_inp = {
"a": np.array(pdf["a"], dtype=np.uint64),
"b": np.array(pdf["b"], dtype=np.float32),
}
# test dataframe input works for 1d tensor specs and input is converted to dict
res = pyfunc_model.predict(pdf)
assert _compare_exact_tensor_dict_input(res, d_inp)
expected_types = dict(
zip(input_schema.input_names(), input_schema.input_types()))
actual_types = {k: v.dtype for k, v in res.items()}
assert expected_types == actual_types
# test that dictionary works too
res = pyfunc_model.predict(d_inp)
assert res == d_inp
expected_types = dict(
zip(input_schema.input_names(), input_schema.input_types()))
actual_types = {k: v.dtype for k, v in res.items()}
assert expected_types == actual_types
def test_schema_enforcement_single_named_tensor_schema():
m = Model()
input_schema = Schema([TensorSpec(np.dtype(np.uint64), (-1, 2), "a")])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
inp = {
"a": np.array([[0, 0], [1, 1]], dtype=np.uint64),
}
# sanity test that dictionary with correct input works
res = pyfunc_model.predict(inp)
assert res == inp
expected_types = dict(
zip(input_schema.input_names(), input_schema.input_types()))
actual_types = {k: v.dtype for k, v in res.items()}
assert expected_types == actual_types
# test single np.ndarray input works and is converted to dictionary
res = pyfunc_model.predict(inp["a"])
assert res == inp
expected_types = dict(
zip(input_schema.input_names(), input_schema.input_types()))
actual_types = {k: v.dtype for k, v in res.items()}
assert expected_types == actual_types
# test list does not work
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict([[0, 0], [1, 1]])
assert "Model is missing inputs ['a']" in str(ex)
def test_model_save_load():
m = Model(artifact_path="some/path",
run_id="123",
flavors={
"flavor1": {
"a": 1,
"b": 2
},
"flavor2": {
"x": 1,
"y": 2
},
},
signature=ModelSignature(
inputs=Schema(
[ColSpec("integer", "x"),
ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")])),
saved_input_example_info={
"x": 1,
"y": 2
})
assert m.get_input_schema() == m.signature.inputs
assert m.get_output_schema() == m.signature.outputs
x = Model(artifact_path="some/other/path", run_id="1234")
assert x.get_input_schema() is None
assert x.get_output_schema() is None
n = Model(artifact_path="some/path",
run_id="123",
flavors={
"flavor1": {
"a": 1,
"b": 2
},
"flavor2": {
"x": 1,
"y": 2
},
},
signature=ModelSignature(
inputs=Schema(
[ColSpec("integer", "x"),
ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")])),
saved_input_example_info={
"x": 1,
"y": 2
})
n.utc_time_created = m.utc_time_created
assert m == n
n.signature = None
assert m != n
with TempDir() as tmp:
m.save(tmp.path("model"))
o = Model.load(tmp.path("model"))
assert m == o
assert m.to_json() == o.to_json()
assert m.to_yaml() == o.to_yaml()
def test_tensor_schema_enforcement_no_col_names():
m = Model()
input_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 3))])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
test_data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
# Can call with numpy array of correct shape
assert np.array_equal(pyfunc_model.predict(test_data), test_data)
# Or can call with a dataframe
assert np.array_equal(pyfunc_model.predict(pd.DataFrame(test_data)),
test_data)
# Can not call with a list
with pytest.raises(
MlflowException,
match=
"This model contains a tensor-based model signature with no input names",
):
pyfunc_model.predict([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
# Can not call with a dict
with pytest.raises(
MlflowException,
match=
"This model contains a tensor-based model signature with no input names",
):
pyfunc_model.predict({"blah": test_data})
# Can not call with a np.ndarray of a wrong shape
with pytest.raises(
MlflowException,
match=re.escape(
"Shape of input (2, 2) does not match expected shape (-1, 3)"),
):
pyfunc_model.predict(np.array([[1.0, 2.0], [4.0, 5.0]]))
# Can not call with a np.ndarray of a wrong type
with pytest.raises(
MlflowException,
match="dtype of input uint32 does not match expected dtype float32"
):
pyfunc_model.predict(test_data.astype(np.uint32))
# Can call with a np.ndarray with more elements along variable axis
test_data2 = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
dtype=np.float32)
assert np.array_equal(pyfunc_model.predict(test_data2), test_data2)
# Can not call with an empty ndarray
with pytest.raises(
MlflowException,
match=re.escape(
"Shape of input () does not match expected shape (-1, 3)")):
pyfunc_model.predict(np.ndarray([]))
def test_schema_enforcement_no_col_names():
class TestModel(object):
@staticmethod
def predict(pdf):
return pdf
m = Model()
input_schema = Schema(
[ColSpec("double"),
ColSpec("double"),
ColSpec("double")])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
test_data = [[1.0, 2.0, 3.0]]
# Can call with just a list
assert pyfunc_model.predict(test_data).equals(pd.DataFrame(test_data))
# Or can call with a DataFrame without column names
assert pyfunc_model.predict(pd.DataFrame(test_data)).equals(
pd.DataFrame(test_data))
# # Or can call with a np.ndarray
assert pyfunc_model.predict(pd.DataFrame(test_data).values).equals(
pd.DataFrame(test_data))
# Or with column names!
pdf = pd.DataFrame(data=test_data, columns=["a", "b", "c"])
assert pyfunc_model.predict(pdf).equals(pdf)
# Must provide the right number of arguments
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict([[1.0, 2.0]])
assert "the provided input only has 2 columns." in str(ex)
# Must provide the right types
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict([[1, 2, 3]])
assert "Can not safely convert int64 to float64" in str(ex)
# Can only provide data type that can be converted to dataframe...
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(set([1, 2, 3]))
assert "Expected input to be DataFrame or list. Found: set" in str(ex)
# 9. dictionaries of str -> list/nparray work
d = {"a": [1.0], "b": [2.0], "c": [3.0]}
assert pyfunc_model.predict(d).equals(pd.DataFrame(d))
def save_model(
model, # vega_saved_model_dir=vega_saved_model_dir,
path,
mlflow_model=None,
conda_env=None,
signature=None,
input_example=None,
style=None,
):
""""""
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path),
DIRECTORY_NOT_EMPTY)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Style-specific Save Logic
current_style = _discovered_styles[f"mlflow_vismod.styles.{style}"]
current_style.Style.save(model, path)
# Saving Conda Environment
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env(style=style)
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
mlflow_model.add_flavor(
FLAVOR_NAME,
# saved_model_dir=MODEL_DIR_SUBPATH,
pickled_model="viz.pkl",
)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow_vismod",
env=conda_env_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(
ft_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = SERIALIZED_MODEL_FILE
model_data_path = os.path.join(path, model_data_subpath)
ft_model.save_model(model_data_path)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
pyfunc.add_to_model(
mlflow_model,
loader_module=__name__,
data=model_data_subpath,
env=conda_env_subpath,
)
mlflow_model.add_flavor(
FLAVOR_NAME, fasttext_version=_get_installed_fasttext_version(), data=model_data_subpath
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(
fastai_learner,
path,
conda_env=None,
code_paths=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
**kwargs,
):
"""
Save a fastai Learner to a path on the local file system.
:param fastai_learner: fastai Learner to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param mlflow_model: MLflow model config this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param kwargs: kwargs to pass to ``Learner.save`` method.
.. code-block:: python
:caption: Example
import os
import mlflow.fastai
# Create a fastai Learner model
model = ...
# Start MLflow session and save model to current working directory
with mlflow.start_run():
model.fit(epochs, learning_rate)
mlflow.fastai.save_model(model, 'model')
# Load saved model for inference
model_uri = "{}/{}".format(os.getcwd(), 'model')
loaded_model = mlflow.fastai.load_model(model_uri)
results = loaded_model.predict(predict_data)
"""
import fastai
from fastai.callback.all import ParamScheduler
from pathlib import Path
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
model_data_subpath = "model.fastai"
model_data_path = os.path.join(path, model_data_subpath)
model_data_path = Path(model_data_path)
os.makedirs(path)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# ParamScheduler currently is not pickable
# hence it is been removed before export and added again after export
cbs = [c for c in fastai_learner.cbs if isinstance(c, ParamScheduler)]
fastai_learner.remove_cbs(cbs)
fastai_learner.export(model_data_path, **kwargs)
fastai_learner.add_cbs(cbs)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.fastai",
data=model_data_subpath,
code=code_dir_subpath,
env=_CONDA_ENV_FILE_NAME,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
fastai_version=fastai.__version__,
data=model_data_subpath,
code=code_dir_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
def save_model(
lgb_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a LightGBM model to a path on the local file system.
:param lgb_model: LightGBM model (an instance of `lightgbm.Booster`_) to be saved.
Note that models that implement the `scikit-learn API`_ are not supported.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this describes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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',
'pip': [
'lightgbm==2.3.0'
]
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
"""
import lightgbm as lgb
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
model_data_subpath = "model.lgb"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save a LightGBM model
lgb_model.save_model(model_data_path)
conda_env, pip_requirements, pip_constraints = (
_process_pip_requirements(
get_default_pip_requirements(), pip_requirements, extra_pip_requirements,
)
if conda_env is None
else _process_conda_env(conda_env)
)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.lightgbm",
data=model_data_subpath,
env=_CONDA_ENV_FILE_NAME,
)
mlflow_model.add_flavor(FLAVOR_NAME, lgb_version=lgb.__version__, data=model_data_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(spacy_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None):
"""
Save a spaCy model to a path on the local file system.
:param spacy_model: spaCy model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this describes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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',
'pip': [
'spacy==2.2.3'
]
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
"""
import spacy
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException(
"Unable to save MLflow model to {path} - path '{path}' "
"already exists".format(path=path))
model_data_subpath = "model.spacy"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(model_data_path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save spacy-model
spacy_model.to_disk(path=model_data_path)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save the pyfunc flavor if at least one text categorizer in spaCy pipeline
if any([
isinstance(pipe_component[1], spacy.pipeline.TextCategorizer)
for pipe_component in spacy_model.pipeline
]):
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.spacy",
data=model_data_subpath,
env=conda_env_subpath)
else:
_logger.warning(
"Generating only the spacy flavor for the provided spacy model. This means the model "
"can be loaded back via `mlflow.spacy.load_model`, but cannot be loaded back using "
"pyfunc APIs like `mlflow.pyfunc.load_model` or via the `mlflow models` CLI commands. "
"MLflow will only generate the pyfunc flavor for spacy models containing a pipeline "
"component that is an instance of spacy.pipeline.TextCategorizer.")
mlflow_model.add_flavor(FLAVOR_NAME,
spacy_version=spacy.__version__,
data=model_data_subpath)
mlflow_model.save(os.path.join(path, "MLmodel"))
def save_model(
pytorch_model,
path,
conda_env=None,
mlflow_model=None,
code_paths=None,
pickle_module=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
requirements_file=None,
extra_files=None,
pip_requirements=None,
extra_pip_requirements=None,
**kwargs,
):
"""
Save a PyTorch model to a path on the local file system.
:param pytorch_model: PyTorch model to be saved. Can be either an eager model (subclass of
``torch.nn.Module``) or scripted model prepared via ``torch.jit.script``
or ``torch.jit.trace``.
The model accept a single ``torch.FloatTensor`` as
input and produce a single output tensor.
If saving an eager model, 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.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param pickle_module: The module that PyTorch should use to serialize ("pickle") the specified
``pytorch_model``. This is passed as the ``pickle_module`` parameter
to ``torch.save()``. By default, this module is also used to
deserialize ("unpickle") the PyTorch model at load time.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param requirements_file:
.. warning::
``requirements_file`` has been deprecated. Please use ``pip_requirements`` instead.
A string containing the path to requirements file. Remote URIs are resolved to absolute
filesystem paths. For example, consider the following ``requirements_file`` string:
.. code-block:: python
requirements_file = "s3://my-bucket/path/to/my_file"
In this case, the ``"my_file"`` requirements file is downloaded from S3. If ``None``,
no requirements file is added to the model.
:param extra_files: A list containing the paths to corresponding extra files. Remote URIs
are resolved to absolute filesystem paths.
For example, consider the following ``extra_files`` list -
extra_files = ["s3://my-bucket/path/to/my_file1",
"s3://my-bucket/path/to/my_file2"]
In this case, the ``"my_file1 & my_file2"`` extra file is downloaded from S3.
If ``None``, no extra files are added to the model.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param kwargs: kwargs to pass to ``torch.save`` method.
.. code-block:: python
:caption: Example
import os
import torch
import mlflow.pytorch
# Class defined here
class LinearNNModel(torch.nn.Module):
...
# Initialize our model, criterion and optimizer
...
# Training loop
...
# Save PyTorch models to current working directory
with mlflow.start_run() as run:
mlflow.pytorch.save_model(model, "model")
# Convert to a scripted model and save it
scripted_pytorch_model = torch.jit.script(model)
mlflow.pytorch.save_model(scripted_pytorch_model, "scripted_model")
# Load each saved model for inference
for model_path in ["model", "scripted_model"]:
model_uri = "{}/{}".format(os.getcwd(), model_path)
loaded_model = mlflow.pytorch.load_model(model_uri)
print("Loaded {}:".format(model_path))
for x in [6.0, 8.0, 12.0, 30.0]:
X = torch.Tensor([[x]])
y_pred = loaded_model(X)
print("predict X: {}, y_pred: {:.2f}".format(x, y_pred.data.item()))
print("--")
.. code-block:: text
:caption: Output
Loaded model:
predict X: 6.0, y_pred: 11.90
predict X: 8.0, y_pred: 15.92
predict X: 12.0, y_pred: 23.96
predict X: 30.0, y_pred: 60.13
--
Loaded scripted_model:
predict X: 6.0, y_pred: 11.90
predict X: 8.0, y_pred: 15.92
predict X: 12.0, y_pred: 23.96
predict X: 30.0, y_pred: 60.13
"""
import torch
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
pickle_module = pickle_module or mlflow_pytorch_pickle_module
if not isinstance(pytorch_model, torch.nn.Module):
raise TypeError("Argument 'pytorch_model' should be a torch.nn.Module")
if code_paths is not None:
if not isinstance(code_paths, list):
raise TypeError("Argument code_paths should be a list, not {}".format(type(code_paths)))
path = os.path.abspath(path)
if os.path.exists(path):
raise RuntimeError("Path '{}' already exists".format(path))
if mlflow_model is None:
mlflow_model = Model()
os.makedirs(path)
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "data"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(model_data_path)
# Persist the pickle module name as a file in the model's `data` directory. This is necessary
# because the `data` directory is the only available parameter to `_load_pyfunc`, and it
# does not contain the MLmodel configuration; therefore, it is not sufficient to place
# the module name in the MLmodel
#
# TODO: Stop persisting this information to the filesystem once we have a mechanism for
# supplying the MLmodel configuration to `mlflow.pytorch._load_pyfunc`
pickle_module_path = os.path.join(model_data_path, _PICKLE_MODULE_INFO_FILE_NAME)
with open(pickle_module_path, "w") as f:
f.write(pickle_module.__name__)
# Save pytorch model
model_path = os.path.join(model_data_path, _SERIALIZED_TORCH_MODEL_FILE_NAME)
if isinstance(pytorch_model, torch.jit.ScriptModule):
torch.jit.ScriptModule.save(pytorch_model, model_path)
else:
torch.save(pytorch_model, model_path, pickle_module=pickle_module, **kwargs)
torchserve_artifacts_config = {}
if extra_files:
torchserve_artifacts_config[_EXTRA_FILES_KEY] = []
if not isinstance(extra_files, list):
raise TypeError("Extra files argument should be a list")
with TempDir() as tmp_extra_files_dir:
for extra_file in extra_files:
_download_artifact_from_uri(
artifact_uri=extra_file, output_path=tmp_extra_files_dir.path()
)
rel_path = posixpath.join(_EXTRA_FILES_KEY, os.path.basename(extra_file))
torchserve_artifacts_config[_EXTRA_FILES_KEY].append({"path": rel_path})
shutil.move(
tmp_extra_files_dir.path(),
posixpath.join(path, _EXTRA_FILES_KEY),
)
if requirements_file:
warnings.warn(
"`requirements_file` has been deprecated. Please use `pip_requirements` instead.",
FutureWarning,
stacklevel=2,
)
if not isinstance(requirements_file, str):
raise TypeError("Path to requirements file should be a string")
with TempDir() as tmp_requirements_dir:
_download_artifact_from_uri(
artifact_uri=requirements_file, output_path=tmp_requirements_dir.path()
)
rel_path = os.path.basename(requirements_file)
torchserve_artifacts_config[_REQUIREMENTS_FILE_KEY] = {"path": rel_path}
shutil.move(tmp_requirements_dir.path(rel_path), path)
if code_paths is not None:
code_dir_subpath = "code"
for code_path in code_paths:
_copy_file_or_tree(src=code_path, dst=path, dst_dir=code_dir_subpath)
else:
code_dir_subpath = None
mlflow_model.add_flavor(
FLAVOR_NAME,
model_data=model_data_subpath,
pytorch_version=str(torch.__version__),
**torchserve_artifacts_config,
)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.pytorch",
data=model_data_subpath,
pickle_module_name=pickle_module.__name__,
code=code_dir_subpath,
env=_CONDA_ENV_FILE_NAME,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
model_data_path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
if not requirements_file:
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
def save_model(
pd_model,
path,
training=False,
conda_env=None,
code_paths=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a paddle model to a path on the local file system. Produces an MLflow Model
containing the following flavors:
- :py:mod:`mlflow.paddle`
- :py:mod:`mlflow.pyfunc`. NOTE: This flavor is only included for paddle models
that define `predict()`, since `predict()` is required for pyfunc model inference.
:param pd_model: paddle model to be saved.
:param path: Local path where the model is to be saved.
:param training: Only valid when saving a model trained using the PaddlePaddle high level API.
If set to True, the saved model supports both re-training and
inference. If set to False, it only supports inference.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
import mlflow.paddle
import paddle
from paddle.nn import Linear
import paddle.nn.functional as F
import numpy as np
import os
import random
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn import preprocessing
def load_data():
# dataset on boston housing prediction
X, y = load_diabetes(return_X_y=True, as_frame=True)
min_max_scaler = preprocessing.MinMaxScaler()
X_min_max = min_max_scaler.fit_transform(X)
X_normalized = preprocessing.scale(X_min_max, with_std=False)
X_train, X_test, y_train, y_test = train_test_split(
X_normalized, y, test_size=0.2, random_state=42)
y_train = y_train.reshape(-1, 1)
y_test = y_test.reshape(-1, 1)
return np.concatenate(
(X_train, y_train), axis=1),np.concatenate((X_test, y_test), axis=1
)
class Regressor(paddle.nn.Layer):
def __init__(self):
super(Regressor, self).__init__()
self.fc = Linear(in_features=13, out_features=1)
@paddle.jit.to_static
def forward(self, inputs):
x = self.fc(inputs)
return x
model = Regressor()
model.train()
training_data, test_data = load_data()
opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
EPOCH_NUM = 10
BATCH_SIZE = 10
for epoch_id in range(EPOCH_NUM):
np.random.shuffle(training_data)
mini_batches = [training_data[k : k + BATCH_SIZE]
for k in range(0, len(training_data), BATCH_SIZE)]
for iter_id, mini_batch in enumerate(mini_batches):
x = np.array(mini_batch[:, :-1]).astype('float32')
y = np.array(mini_batch[:, -1:]).astype('float32')
house_features = paddle.to_tensor(x)
prices = paddle.to_tensor(y)
predicts = model(house_features)
loss = F.square_error_cost(predicts, label=prices)
avg_loss = paddle.mean(loss)
if iter_id%20==0:
print("epoch: {}, iter: {}, loss is: {}".format(
epoch_id, iter_id, avg_loss.numpy()))
avg_loss.backward()
opt.step()
opt.clear_grad()
mlflow.log_param('learning_rate', 0.01)
mlflow.paddle.log_model(model, "model")
sk_path_dir = './test-out'
mlflow.paddle.save_model(model, sk_path_dir)
print("Model saved in run %s" % mlflow.active_run().info.run_uuid)
"""
import paddle
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
_validate_and_prepare_target_save_path(path)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "model"
output_path = os.path.join(path, model_data_subpath)
if isinstance(pd_model, paddle.Model):
pd_model.save(output_path, training=training)
else:
paddle.jit.save(pd_model, output_path)
# `PyFuncModel` only works for paddle models that define `predict()`.
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.paddle",
model_path=model_data_subpath,
env=_CONDA_ENV_FILE_NAME,
code=code_dir_subpath,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
pickled_model=model_data_subpath,
paddle_version=paddle.__version__,
code=code_dir_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
_PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))
def save_model(
pytorch_model,
path,
conda_env=None,
mlflow_model=None,
code_paths=None,
pickle_module=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
**kwargs
):
"""
Save a PyTorch model to a path on the local file system.
:param 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.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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'
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param pickle_module: The module that PyTorch should use to serialize ("pickle") the specified
``pytorch_model``. This is passed as the ``pickle_module`` parameter
to ``torch.save()``. By default, this module is also used to
deserialize ("unpickle") the PyTorch model at load time.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param kwargs: kwargs to pass to ``torch.save`` method.
.. code-block:: python
:caption: Example
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 = pytorch_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)
"""
import torch
pickle_module = pickle_module or mlflow_pytorch_pickle_module
if not isinstance(pytorch_model, torch.nn.Module):
raise TypeError("Argument 'pytorch_model' should be a torch.nn.Module")
if code_paths is not None:
if not isinstance(code_paths, list):
raise TypeError("Argument code_paths should be a list, not {}".format(type(code_paths)))
path = os.path.abspath(path)
if os.path.exists(path):
raise RuntimeError("Path '{}' already exists".format(path))
if mlflow_model is None:
mlflow_model = Model()
os.makedirs(path)
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "data"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(model_data_path)
# Persist the pickle module name as a file in the model's `data` directory. This is necessary
# because the `data` directory is the only available parameter to `_load_pyfunc`, and it
# does not contain the MLmodel configuration; therefore, it is not sufficient to place
# the module name in the MLmodel
#
# TODO: Stop persisting this information to the filesystem once we have a mechanism for
# supplying the MLmodel configuration to `mlflow.pytorch._load_pyfunc`
pickle_module_path = os.path.join(model_data_path, _PICKLE_MODULE_INFO_FILE_NAME)
with open(pickle_module_path, "w") as f:
f.write(pickle_module.__name__)
# Save pytorch model
model_path = os.path.join(model_data_path, _SERIALIZED_TORCH_MODEL_FILE_NAME)
torch.save(pytorch_model, model_path, pickle_module=pickle_module, **kwargs)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
if code_paths is not None:
code_dir_subpath = "code"
for code_path in code_paths:
_copy_file_or_tree(src=code_path, dst=path, dst_dir=code_dir_subpath)
else:
code_dir_subpath = None
mlflow_model.add_flavor(
FLAVOR_NAME, model_data=model_data_subpath, pytorch_version=torch.__version__
)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.pytorch",
data=model_data_subpath,
pickle_module_name=pickle_module.__name__,
code=code_dir_subpath,
env=conda_env_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(cb_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
**kwargs):
"""
Save a CatBoost model to a path on the local file system.
:param cb_model: CatBoost model (an instance of `CatBoost`_, `CatBoostClassifier`_,
or `CatBoostRegressor`_) to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param kwargs: kwargs to pass to `CatBoost.save_model`_ method.
"""
import catboost as cb
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_path = os.path.join(path, _MODEL_BINARY_FILE_NAME)
cb_model.save_model(model_data_path, **kwargs)
model_bin_kwargs = {_MODEL_BINARY_KEY: _MODEL_BINARY_FILE_NAME}
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.catboost",
env=_CONDA_ENV_FILE_NAME,
**model_bin_kwargs,
)
flavor_conf = {
_MODEL_TYPE_KEY: cb_model.__class__.__name__,
_SAVE_FORMAT_KEY: kwargs.get("format", "cbm"),
**model_bin_kwargs,
}
mlflow_model.add_flavor(
FLAVOR_NAME,
catboost_version=cb.__version__,
**flavor_conf,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
def save_model(
tf_saved_model_dir,
tf_meta_graph_tags,
tf_signature_def_key,
path,
mlflow_model=None,
conda_env=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
"""
Save a *serialized* collection of TensorFlow graphs and variables as an MLflow model
to a local path. This method operates on TensorFlow variables and graphs that have been
serialized in TensorFlow's ``SavedModel`` format. For more information about ``SavedModel``
format, see the TensorFlow documentation:
https://www.tensorflow.org/guide/saved_model#save_and_restore_models.
:param tf_saved_model_dir: Path to the directory containing serialized TensorFlow variables and
graphs in ``SavedModel`` format.
:param tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
serialized ``SavedModel`` object. For more information, see the
``tags`` parameter of the
``tf.saved_model.builder.savedmodelbuilder`` method.
:param tf_signature_def_key: A string identifying the input/output signature associated with the
model. This is a key within the serialized ``savedmodel``
signature definition mapping. For more information, see the
``signature_def_map`` parameter of the
``tf.saved_model.builder.savedmodelbuilder`` method.
:param path: Local path where the MLflow model is to be saved.
:param mlflow_model: MLflow model configuration to which to add the ``tensorflow`` flavor.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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',
'tensorflow=1.8.0'
]
}
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
"""
_logger.info(
"Validating the specified TensorFlow model by attempting to load it in a new TensorFlow"
" graph...")
_validate_saved_model(
tf_saved_model_dir=tf_saved_model_dir,
tf_meta_graph_tags=tf_meta_graph_tags,
tf_signature_def_key=tf_signature_def_key,
)
_logger.info("Validation succeeded!")
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path),
DIRECTORY_NOT_EMPTY)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
root_relative_path = _copy_file_or_tree(src=tf_saved_model_dir,
dst=path,
dst_dir=None)
model_dir_subpath = "tfmodel"
shutil.move(os.path.join(path, root_relative_path),
os.path.join(path, model_dir_subpath))
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
mlflow_model.add_flavor(
FLAVOR_NAME,
saved_model_dir=model_dir_subpath,
meta_graph_tags=tf_meta_graph_tags,
signature_def_key=tf_signature_def_key,
)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.tensorflow",
env=conda_env_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(fastai_learner,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
**kwargs):
"""
Save a fastai Learner to a path on the local file system.
:param fastai_learner: fastai Learner to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this describes the environment
this model should be run in. At minimum, it should specify the
dependencies contained in :func:`get_default_conda_env()`. If
``None``, the default :func:`get_default_conda_env()` environment is
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',
'fastai=1.0.60',
]
}
:param mlflow_model: MLflow model config this flavor is being added to.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param kwargs: kwargs to pass to ``Learner.save`` method.
"""
import fastai
from pathlib import Path
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
model_data_subpath = "model.fastai"
model_data_path = os.path.join(path, model_data_subpath)
model_data_path = Path(model_data_path)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save an Learner
fastai_learner.export(model_data_path, **kwargs)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.fastai",
data=model_data_subpath,
env=conda_env_subpath)
mlflow_model.add_flavor(FLAVOR_NAME,
fastai_version=fastai.__version__,
data=model_data_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(cb_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
**kwargs):
"""
Save a CatBoost model to a path on the local file system.
:param cb_model: CatBoost model (an instance of `CatBoost`_, `CatBoostClassifier`_,
or `CatBoostRegressor`_) to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this describes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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',
'pip': [
'catboost==0.19.1'
]
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param kwargs: kwargs to pass to `CatBoost.save_model`_ method.
"""
import catboost as cb
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_path = os.path.join(path, _MODEL_BINARY_FILE_NAME)
cb_model.save_model(model_data_path, **kwargs)
conda_env, pip_requirements, pip_constraints = (_process_pip_requirements(
get_default_pip_requirements(),
pip_requirements,
extra_pip_requirements,
) if conda_env is None else _process_conda_env(conda_env))
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
model_bin_kwargs = {_MODEL_BINARY_KEY: _MODEL_BINARY_FILE_NAME}
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.catboost",
env=_CONDA_ENV_FILE_NAME,
**model_bin_kwargs,
)
flavor_conf = {
_MODEL_TYPE_KEY: cb_model.__class__.__name__,
_SAVE_FORMAT_KEY: kwargs.get("format", "cbm"),
**model_bin_kwargs,
}
mlflow_model.add_flavor(
FLAVOR_NAME,
catboost_version=cb.__version__,
**flavor_conf,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(keras_model,
path,
conda_env=None,
mlflow_model=None,
custom_objects=None,
keras_module=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
**kwargs):
"""
Save a Keras model to a path on the local file system.
:param keras_model: Keras model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the
dependencies contained in :func:`get_default_conda_env()`. If
``None``, the default :func:`get_default_conda_env()` environment is
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',
'keras=2.2.4',
'tensorflow=1.8.0'
]
}
:param mlflow_model: MLflow model config this flavor is being added to.
:param custom_objects: A Keras ``custom_objects`` dictionary mapping names (strings) to
custom classes or functions associated with the Keras model. MLflow saves
these custom layers using CloudPickle and restores them automatically
when the model is loaded with :py:func:`mlflow.keras.load_model` and
:py:func:`mlflow.pyfunc.load_model`.
:param keras_module: Keras module to be used to save / load the model
(``keras`` or ``tf.keras``). If not provided, MLflow will
attempt to infer the Keras module based on the given model.
:param kwargs: kwargs to pass to ``keras_model.save`` method.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
.. code-block:: python
:caption: Example
import mlflow
# Build, compile, and train your model
keras_model = ...
keras_model_path = ...
keras_model.compile(optimizer="rmsprop", loss="mse", metrics=["accuracy"])
results = keras_model.fit(
x_train, y_train, epochs=20, batch_size = 128, validation_data=(x_val, y_val))
# Save the model as an MLflow Model
mlflow.keras.save_model(keras_model, keras_model_path)
"""
if keras_module is None:
def _is_plain_keras(model):
try:
# NB: Network is the first parent with save method
import keras.engine.network
return isinstance(model, keras.engine.network.Network)
except ImportError:
return False
def _is_tf_keras(model):
try:
# NB: Network is not exposed in tf.keras, we check for Model instead.
import tensorflow.keras.models
return isinstance(model, tensorflow.keras.models.Model)
except ImportError:
return False
if _is_plain_keras(keras_model):
keras_module = importlib.import_module("keras")
elif _is_tf_keras(keras_model):
keras_module = importlib.import_module("tensorflow.keras")
else:
raise MlflowException(
"Unable to infer keras module from the model, please specify "
"which keras module ('keras' or 'tensorflow.keras') is to be "
"used to save and load the model.")
elif type(keras_module) == str:
keras_module = importlib.import_module(keras_module)
# check if path exists
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
# construct new data folder in existing path
data_subpath = "data"
data_path = os.path.join(path, data_subpath)
os.makedirs(data_path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# save custom objects if there are custom objects
if custom_objects is not None:
_save_custom_objects(data_path, custom_objects)
# save keras module spec to path/data/keras_module.txt
with open(os.path.join(data_path, _KERAS_MODULE_SPEC_PATH), "w") as f:
f.write(keras_module.__name__)
# save keras model to path/data/model.h5
model_subpath = os.path.join(data_subpath, _MODEL_SAVE_PATH)
model_path = os.path.join(path, model_subpath)
if path.startswith('/dbfs/'):
# The Databricks Filesystem uses a FUSE implementation that does not support
# random writes. It causes an error.
with tempfile.NamedTemporaryFile(suffix='.h5') as f:
keras_model.save(f.name, **kwargs)
f.flush() # force flush the data
shutil.copyfile(src=f.name, dst=model_path)
else:
keras_model.save(model_path, **kwargs)
# update flavor info to mlflow_model
mlflow_model.add_flavor(FLAVOR_NAME,
keras_module=keras_module.__name__,
keras_version=keras_module.__version__,
data=data_subpath)
# save conda.yaml info to path/conda.yml
if conda_env is None:
conda_env = get_default_conda_env(include_cloudpickle=custom_objects
is not None,
keras_module=keras_module)
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, _CONDA_ENV_SUBPATH), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# append loader_module, data and env data to mlflow_model
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.keras",
data=data_subpath,
env=_CONDA_ENV_SUBPATH)
# save mlflow_model to path/MLmodel
mlflow_model.save(os.path.join(path, "MLmodel"))
def save_model(
gluon_model,
path,
mlflow_model=None,
conda_env=None,
code_paths=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a Gluon model to a path on the local file system.
:param gluon_model: Gluon model to be saved. Must be already hybridized.
:param path: Local path where the model is to be saved.
:param mlflow_model: MLflow model config this flavor is being added to.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
from mxnet.gluon import Trainer
from mxnet.gluon.contrib import estimator
from mxnet.gluon.loss import SoftmaxCrossEntropyLoss
from mxnet.gluon.nn import HybridSequential
from mxnet.metric import Accuracy
import mlflow
# Build, compile, and train your model
gluon_model_path = ...
net = HybridSequential()
with net.name_scope():
...
net.hybridize()
net.collect_params().initialize()
softmax_loss = SoftmaxCrossEntropyLoss()
trainer = Trainer(net.collect_params())
est = estimator.Estimator(net=net, loss=softmax_loss, metrics=Accuracy(), trainer=trainer)
est.fit(train_data=train_data, epochs=100, val_data=validation_data)
# Save the model as an MLflow Model
mlflow.gluon.save_model(net, gluon_model_path)
"""
import mxnet as mx
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
path = os.path.abspath(path)
_validate_and_prepare_target_save_path(path)
data_subpath = "data"
data_path = os.path.join(path, data_subpath)
os.makedirs(data_path)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# The epoch argument of the export method does not play any role in selecting
# a specific epoch's parameters, and is there only for display purposes.
gluon_model.export(os.path.join(data_path, _MODEL_SAVE_PATH))
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.gluon",
env=_CONDA_ENV_FILE_NAME,
code=code_dir_subpath)
mlflow_model.add_flavor(FLAVOR_NAME,
mxnet_version=mx.__version__,
code=code_dir_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
_PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))
def test_column_schema_enforcement():
m = Model()
input_schema = Schema([
ColSpec("integer", "a"),
ColSpec("long", "b"),
ColSpec("float", "c"),
ColSpec("double", "d"),
ColSpec("boolean", "e"),
ColSpec("string", "g"),
ColSpec("binary", "f"),
ColSpec("datetime", "h"),
])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
pdf = pd.DataFrame(
data=[[
1, 2, 3, 4, True, "x",
bytes([1]), "2021-01-01 00:00:00.1234567"
]],
columns=["b", "d", "a", "c", "e", "g", "f", "h"],
dtype=np.object,
)
pdf["a"] = pdf["a"].astype(np.int32)
pdf["b"] = pdf["b"].astype(np.int64)
pdf["c"] = pdf["c"].astype(np.float32)
pdf["d"] = pdf["d"].astype(np.float64)
pdf["h"] = pdf["h"].astype(np.datetime64)
# test that missing column raises
with pytest.raises(MlflowException) as ex:
res = pyfunc_model.predict(pdf[["b", "d", "a", "e", "g", "f", "h"]])
assert "Model is missing inputs" in str(ex)
# test that extra column is ignored
pdf["x"] = 1
# test that columns are reordered, extra column is ignored
res = pyfunc_model.predict(pdf)
assert all((res == pdf[input_schema.input_names()]).all())
expected_types = dict(
zip(input_schema.input_names(), input_schema.pandas_types()))
# MLflow datetime type in input_schema does not encode precision, so add it for assertions
expected_types["h"] = np.dtype("datetime64[ns]")
actual_types = res.dtypes.to_dict()
assert expected_types == actual_types
# Test conversions
# 1. long -> integer raises
pdf["a"] = pdf["a"].astype(np.int64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["a"] = pdf["a"].astype(np.int32)
# 2. integer -> long works
pdf["b"] = pdf["b"].astype(np.int32)
res = pyfunc_model.predict(pdf)
assert all((res == pdf[input_schema.input_names()]).all())
assert res.dtypes.to_dict() == expected_types
pdf["b"] = pdf["b"].astype(np.int64)
# 3. unsigned int -> long works
pdf["b"] = pdf["b"].astype(np.uint32)
res = pyfunc_model.predict(pdf)
assert all((res == pdf[input_schema.input_names()]).all())
assert res.dtypes.to_dict() == expected_types
pdf["b"] = pdf["b"].astype(np.int64)
# 4. unsigned int -> int raises
pdf["a"] = pdf["a"].astype(np.uint32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["a"] = pdf["a"].astype(np.int32)
# 5. double -> float raises
pdf["c"] = pdf["c"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["c"] = pdf["c"].astype(np.float32)
# 6. float -> double works, double -> float does not
pdf["d"] = pdf["d"].astype(np.float32)
res = pyfunc_model.predict(pdf)
assert res.dtypes.to_dict() == expected_types
assert "Incompatible input types" in str(ex)
pdf["d"] = pdf["d"].astype(np.float64)
pdf["c"] = pdf["c"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["c"] = pdf["c"].astype(np.float32)
# 7. int -> float raises
pdf["c"] = pdf["c"].astype(np.int32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["c"] = pdf["c"].astype(np.float32)
# 8. int -> double works
pdf["d"] = pdf["d"].astype(np.int32)
pyfunc_model.predict(pdf)
assert all((res == pdf[input_schema.input_names()]).all())
assert res.dtypes.to_dict() == expected_types
# 9. long -> double raises
pdf["d"] = pdf["d"].astype(np.int64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["d"] = pdf["d"].astype(np.float64)
# 10. any float -> any int raises
pdf["a"] = pdf["a"].astype(np.float32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
# 10. any float -> any int raises
pdf["a"] = pdf["a"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["a"] = pdf["a"].astype(np.int32)
pdf["b"] = pdf["b"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["b"] = pdf["b"].astype(np.int64)
pdf["b"] = pdf["b"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
pdf["b"] = pdf["b"].astype(np.int64)
assert "Incompatible input types" in str(ex)
# 11. objects work
pdf["b"] = pdf["b"].astype(np.object)
pdf["d"] = pdf["d"].astype(np.object)
pdf["e"] = pdf["e"].astype(np.object)
pdf["f"] = pdf["f"].astype(np.object)
pdf["g"] = pdf["g"].astype(np.object)
res = pyfunc_model.predict(pdf)
assert res.dtypes.to_dict() == expected_types
# 12. datetime64[D] (date only) -> datetime64[x] works
pdf["h"] = pdf["h"].astype("datetime64[D]")
res = pyfunc_model.predict(pdf)
assert res.dtypes.to_dict() == expected_types
pdf["h"] = pdf["h"].astype("datetime64[s]")
# 13. np.ndarrays can be converted to dataframe but have no columns
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf.values)
assert "Model is missing inputs" in str(ex)
# 14. dictionaries of str -> list/nparray work
arr = np.array([1, 2, 3])
d = {
"a":
arr.astype("int32"),
"b":
arr.astype("int64"),
"c":
arr.astype("float32"),
"d":
arr.astype("float64"),
"e": [True, False, True],
"g": ["a", "b", "c"],
"f": [bytes(0), bytes(1), bytes(1)],
"h":
np.array(["2020-01-01", "2020-02-02", "2020-03-03"],
dtype=np.datetime64),
}
res = pyfunc_model.predict(d)
assert res.dtypes.to_dict() == expected_types
# 15. dictionaries of str -> list[list] fail
d = {
"a": [arr.astype("int32")],
"b": [arr.astype("int64")],
"c": [arr.astype("float32")],
"d": [arr.astype("float64")],
"e": [[True, False, True]],
"g": [["a", "b", "c"]],
"f": [[bytes(0), bytes(1), bytes(1)]],
"h": [
np.array(["2020-01-01", "2020-02-02", "2020-03-03"],
dtype=np.datetime64)
],
}
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(d)
assert "Incompatible input types" in str(ex)
# 16. conversion to dataframe fails
d = {
"a": [1],
"b": [1, 2],
"c": [1, 2, 3],
}
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(d)
assert "This model contains a column-based signature, which suggests a DataFrame input." in str(
ex)
def test_tensor_multi_named_schema_enforcement():
m = Model()
input_schema = Schema([
TensorSpec(np.dtype(np.uint64), (-1, 5), "a"),
TensorSpec(np.dtype(np.short), (-1, 2), "b"),
TensorSpec(np.dtype(np.float32), (2, -1, 2), "c"),
])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
inp = {
"a": np.array([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], dtype=np.uint64),
"b": np.array([[0, 0], [1, 1], [2, 2]], dtype=np.short),
"c": np.array([[[0, 0], [1, 1]], [[2, 2], [3, 3]]], dtype=np.float32),
}
# test that missing column raises
inp1 = {k: v for k, v in inp.items()}
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(inp1.pop("b"))
assert "Model is missing inputs" in str(ex)
# test that extra column is ignored
inp2 = {k: v for k, v in inp.items()}
inp2["x"] = 1
# test that extra column is removed
res = pyfunc_model.predict(inp2)
assert res == {k: v for k, v in inp.items() if k in {"a", "b", "c"}}
expected_types = dict(
zip(input_schema.input_names(), input_schema.input_types()))
actual_types = {k: v.dtype for k, v in res.items()}
assert expected_types == actual_types
# test that variable axes are supported
inp3 = {
"a":
np.array([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 2, 2, 2]],
dtype=np.uint64),
"b":
np.array([[0, 0], [1, 1]], dtype=np.short),
"c":
np.array([[[0, 0]], [[2, 2]]], dtype=np.float32),
}
res = pyfunc_model.predict(inp3)
assert _compare_exact_tensor_dict_input(res, inp3)
expected_types = dict(
zip(input_schema.input_names(), input_schema.input_types()))
actual_types = {k: v.dtype for k, v in res.items()}
assert expected_types == actual_types
# test that type casting is not supported
inp4 = {k: v for k, v in inp.items()}
inp4["a"] = inp4["a"].astype(np.int32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(inp4)
assert "dtype of input int32 does not match expected dtype uint64" in str(
ex)
# test wrong shape
inp5 = {
"a": np.array([[0, 0, 0, 0]], dtype=np.uint),
"b": np.array([[0, 0], [1, 1]], dtype=np.short),
"c": np.array([[[0, 0]]], dtype=np.float32),
}
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(inp5)
assert "Shape of input (1, 4) does not match expected shape (-1, 5)" in str(
ex)
# test non-dictionary input
inp6 = [
np.array([[0, 0, 0, 0, 0]], dtype=np.uint64),
np.array([[0, 0], [1, 1]], dtype=np.short),
np.array([[[0, 0]]], dtype=np.float32),
]
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(inp6)
assert "Model is missing inputs ['a', 'b', 'c']." in str(ex)
# test empty ndarray does not work
inp7 = {k: v for k, v in inp.items()}
inp7["a"] = np.array([])
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(inp7)
assert "Shape of input (0,) does not match expected shape" in str(ex)
# test dictionary of str -> list does not work
inp8 = {k: list(v) for k, v in inp.items()}
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(inp8)
assert "This model contains a tensor-based model signature with input names" in str(
ex)
assert (
"suggests a dictionary input mapping input name to a numpy array, but a dict"
" with value type <class 'list'> was found") in str(ex)
# test dataframe input fails at shape enforcement
pdf = pd.DataFrame(
data=[[1, 2, 3]],
columns=["a", "b", "c"],
)
pdf["a"] = pdf["a"].astype(np.uint64)
pdf["b"] = pdf["b"].astype(np.short)
pdf["c"] = pdf["c"].astype(np.float32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Shape of input (1,) does not match expected shape (-1, 5)" in str(
ex)
def save_model(
lgb_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a LightGBM model to a path on the local file system.
:param lgb_model: LightGBM model (an instance of `lightgbm.Booster`_) to be saved.
Note that models that implement the `scikit-learn API`_ are not supported.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
"""
import lightgbm as lgb
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
model_data_subpath = "model.lgb"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save a LightGBM model
lgb_model.save_model(model_data_path)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.lightgbm",
data=model_data_subpath,
env=_CONDA_ENV_FILE_NAME,
)
mlflow_model.add_flavor(FLAVOR_NAME, lgb_version=lgb.__version__, data=model_data_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path, FLAVOR_NAME, fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs, pip_requirements, extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
def test_schema_enforcement():
class TestModel(object):
@staticmethod
def predict(pdf):
return pdf
m = Model()
input_schema = Schema([
ColSpec("integer", "a"),
ColSpec("long", "b"),
ColSpec("float", "c"),
ColSpec("double", "d"),
ColSpec("boolean", "e"),
ColSpec("string", "g"),
ColSpec("binary", "f"),
])
m.signature = ModelSignature(inputs=input_schema)
pyfunc_model = PyFuncModel(model_meta=m, model_impl=TestModel())
pdf = pd.DataFrame(
data=[[1, 2, 3, 4, True, "x", bytes([1])]],
columns=["b", "d", "a", "c", "e", "g", "f"],
dtype=np.object,
)
pdf["a"] = pdf["a"].astype(np.int32)
pdf["b"] = pdf["b"].astype(np.int64)
pdf["c"] = pdf["c"].astype(np.float32)
pdf["d"] = pdf["d"].astype(np.float64)
# test that missing column raises
with pytest.raises(MlflowException) as ex:
res = pyfunc_model.predict(pdf[["b", "d", "a", "e", "g", "f"]])
assert "Model input is missing columns" in str(ex)
# test that extra column is ignored
pdf["x"] = 1
# test that columns are reordered, extra column is ignored
res = pyfunc_model.predict(pdf)
assert all((res == pdf[input_schema.column_names()]).all())
expected_types = dict(
zip(input_schema.column_names(), input_schema.pandas_types()))
actual_types = res.dtypes.to_dict()
assert expected_types == actual_types
# Test conversions
# 1. long -> integer raises
pdf["a"] = pdf["a"].astype(np.int64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["a"] = pdf["a"].astype(np.int32)
# 2. integer -> long works
pdf["b"] = pdf["b"].astype(np.int32)
res = pyfunc_model.predict(pdf)
assert all((res == pdf[input_schema.column_names()]).all())
assert res.dtypes.to_dict() == expected_types
pdf["b"] = pdf["b"].astype(np.int64)
# 3. double -> float raises
pdf["c"] = pdf["c"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["c"] = pdf["c"].astype(np.float32)
# 4. float -> double works
pdf["d"] = pdf["d"].astype(np.float32)
res = pyfunc_model.predict(pdf)
assert res.dtypes.to_dict() == expected_types
assert "Incompatible input types" in str(ex)
pdf["d"] = pdf["d"].astype(np.int64)
# 5. floats -> ints raises
pdf["c"] = pdf["c"].astype(np.int32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["c"] = pdf["c"].astype(np.float32)
pdf["d"] = pdf["d"].astype(np.int64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["d"] = pdf["d"].astype(np.float64)
# 6. ints -> floats raises
pdf["a"] = pdf["a"].astype(np.float32)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
assert "Incompatible input types" in str(ex)
pdf["a"] = pdf["a"].astype(np.int32)
pdf["b"] = pdf["b"].astype(np.float64)
with pytest.raises(MlflowException) as ex:
pyfunc_model.predict(pdf)
pdf["b"] = pdf["b"].astype(np.int64)
assert "Incompatible input types" in str(ex)
# 7. objects work
pdf["b"] = pdf["b"].astype(np.object)
pdf["d"] = pdf["d"].astype(np.object)
pdf["e"] = pdf["e"].astype(np.object)
pdf["f"] = pdf["f"].astype(np.object)
pdf["g"] = pdf["g"].astype(np.object)
res = pyfunc_model.predict(pdf)
assert res.dtypes.to_dict() == expected_types
def save_model(
gluon_model,
path,
mlflow_model=None,
conda_env=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
"""
Save a Gluon model to a path on the local file system.
:param gluon_model: Gluon model to be saved. Must be already hybridized.
:param path: Local path where the model is to be saved.
:param mlflow_model: MLflow model config this flavor is being added to.
:param 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 :func:`get_default_conda_env()`. If ``None``, the default
:func:`mlflow.gluon.get_default_conda_env()` environment is 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',
'mxnet=1.5.0'
]
}
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
.. code-block:: python
:caption: Example
from mxnet.gluon import Trainer
from mxnet.gluon.contrib import estimator
from mxnet.gluon.loss import SoftmaxCrossEntropyLoss
from mxnet.gluon.nn import HybridSequential
from mxnet.metric import Accuracy
import mlflow
# Build, compile, and train your model
gluon_model_path = ...
net = HybridSequential()
with net.name_scope():
...
net.hybridize()
net.collect_params().initialize()
softmax_loss = SoftmaxCrossEntropyLoss()
trainer = Trainer(net.collect_params())
est = estimator.Estimator(net=net, loss=softmax_loss, metrics=Accuracy(), trainer=trainer)
est.fit(train_data=train_data, epochs=100, val_data=validation_data)
# Save the model as an MLflow Model
mlflow.gluon.save_model(net, gluon_model_path)
"""
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
data_subpath = "data"
data_path = os.path.join(path, data_subpath)
os.makedirs(data_path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# The epoch argument of the export method does not play any role in selecting
# a specific epoch's paramaters, and is there only for display purposes.
gluon_model.export(os.path.join(data_path, _MODEL_SAVE_PATH))
with open(os.path.join(path, "architecture.txt"), "w") as fp:
fp.write(str(gluon_model))
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.gluon",
env=conda_env_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(
spark_model,
sample_input,
path,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
"""
Save a Spark MLlib PipelineModel in MLeap format at a local path.
The saved model will have the MLeap flavor.
NOTE:
You cannot load the MLeap model flavor in Python; you must download it using the
Java API method ``downloadArtifacts(String runId)`` and load the model
using the method ``MLeapLoader.loadPipeline(String modelRootPath)``.
:param spark_model: Spark PipelineModel to be saved. This model must be MLeap-compatible and
cannot contain any custom transformers.
:param sample_input: Sample PySpark DataFrame input that the model can evaluate. This is
required by MLeap for data schema inference.
:param path: Local path where the model is to be saved.
:param mlflow_model: :py:mod:`mlflow.models.Model` to which this flavor is being added.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset) and valid
model output (e.g. model predictions generated on the training dataset),
for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
signature = infer_signature(train, model.predict(train))
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
"""
if mlflow_model is None:
mlflow_model = Model()
add_to_model(mlflow_model=mlflow_model,
path=path,
spark_model=spark_model,
sample_input=sample_input)
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(path,
loader_module=None,
data_path=None,
code_path=None,
conda_env=None,
mlflow_model=None,
python_model=None,
artifacts=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
**kwargs):
"""
save_model(path, loader_module=None, data_path=None, code_path=None, conda_env=None,\
mlflow_model=Model(), python_model=None, artifacts=None)
Save a Pyfunc model with custom inference logic and optional data dependencies to a path on the
local filesystem.
For information about the workflows that this method supports, please see :ref:`"workflows for
creating custom pyfunc models" <pyfunc-create-custom-workflows>` and
:ref:`"which workflow is right for my use case?" <pyfunc-create-custom-selecting-workflow>`.
Note that the parameters for the second workflow: ``loader_module``, ``data_path`` and the
parameters for the first workflow: ``python_model``, ``artifacts``, cannot be
specified together.
:param path: The path to which to save the Python model.
:param loader_module: The name of the Python module that is used to load the model
from ``data_path``. This module must define a method with the prototype
``_load_pyfunc(data_path)``. If not ``None``, this module and its
dependencies must be included in one of the following locations:
- The MLflow library.
- 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_path`` parameter.
:param data_path: Path to a file or directory containing model data.
:param code_path: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path before the model is loaded.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. This decsribes the environment this model should
be run in. If ``python_model`` is not ``None``, the Conda environment must
at least specify the dependencies contained in
:func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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',
'cloudpickle==0.5.8'
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` configuration to which to add the
**python_function** flavor.
:param python_model: An instance of a subclass of :class:`~PythonModel`. This class is
serialized using the CloudPickle library. Any dependencies of the class
should be included in one of the following locations:
- The MLflow library.
- 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_path`` parameter.
Note: If the class is imported from another module, as opposed to being
defined in the ``__main__`` scope, the defining module should also be
included in one of the listed locations.
:param artifacts: A dictionary containing ``<name, artifact_uri>`` entries. Remote artifact URIs
are resolved to absolute filesystem paths, producing a dictionary of
``<name, absolute_path>`` entries. ``python_model`` can reference these
resolved entries as the ``artifacts`` property of the ``context`` parameter
in :func:`PythonModel.load_context() <mlflow.pyfunc.PythonModel.load_context>`
and :func:`PythonModel.predict() <mlflow.pyfunc.PythonModel.predict>`.
For example, consider the following ``artifacts`` dictionary::
{
"my_file": "s3://my-bucket/path/to/my/file"
}
In this case, the ``"my_file"`` artifact is downloaded from S3. The
``python_model`` can then refer to ``"my_file"`` as an absolute filesystem
path via ``context.artifacts["my_file"]``.
If ``None``, no artifacts are added to the model.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
"""
mlflow_model = kwargs.pop("model", mlflow_model)
if len(kwargs) > 0:
raise TypeError(
"save_model() got unexpected keyword arguments: {}".format(kwargs))
if code_path is not None:
if not isinstance(code_path, list):
raise TypeError(
"Argument code_path should be a list, not {}".format(
type(code_path)))
first_argument_set = {
"loader_module": loader_module,
"data_path": data_path,
}
second_argument_set = {
"artifacts": artifacts,
"python_model": python_model,
}
first_argument_set_specified = any(
[item is not None for item in first_argument_set.values()])
second_argument_set_specified = any(
[item is not None for item in second_argument_set.values()])
if first_argument_set_specified and second_argument_set_specified:
raise MlflowException(
message=
("The following sets of parameters cannot be specified together: {first_set_keys}"
" and {second_set_keys}. All parameters in one set must be `None`. Instead, found"
" the following values: {first_set_entries} and {second_set_entries}"
.format(
first_set_keys=first_argument_set.keys(),
second_set_keys=second_argument_set.keys(),
first_set_entries=first_argument_set,
second_set_entries=second_argument_set,
)),
error_code=INVALID_PARAMETER_VALUE,
)
elif (loader_module is None) and (python_model is None):
msg = (
"Either `loader_module` or `python_model` must be specified. A `loader_module` "
"should be a python module. A `python_model` should be a subclass of PythonModel"
)
raise MlflowException(message=msg, error_code=INVALID_PARAMETER_VALUE)
if os.path.exists(path):
raise MlflowException(message="Path '{}' already exists".format(path),
error_code=RESOURCE_ALREADY_EXISTS)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
if first_argument_set_specified:
return _save_model_with_loader_module_and_data_path(
path=path,
loader_module=loader_module,
data_path=data_path,
code_paths=code_path,
conda_env=conda_env,
mlflow_model=mlflow_model,
)
elif second_argument_set_specified:
return mlflow.pyfunc.model._save_model_with_class_artifacts_params(
path=path,
python_model=python_model,
artifacts=artifacts,
conda_env=conda_env,
code_paths=code_path,
mlflow_model=mlflow_model,
)
def save_explainer(
explainer,
path,
serialize_model_using_mlflow=True,
conda_env=None,
code_paths=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a SHAP explainer to a path on the local file system. Produces an MLflow Model
containing the following flavors:
- :py:mod:`mlflow.shap`
- :py:mod:`mlflow.pyfunc`
:param explainer: SHAP explainer to be saved.
:param path: Local path where the explainer is to be saved.
:param serialize_model_using_mlflow: When set to True, MLflow will extract the underlying
model and serialize it as an MLmodel, otherwise it
uses SHAP's internal serialization. Defaults to True.
Currently MLflow serialization is only supported for
models of 'sklearn' or 'pytorch' flavors.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
"""
import shap
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
_validate_and_prepare_target_save_path(path)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
underlying_model_flavor = None
underlying_model_path = None
serializable_by_mlflow = False
# saving the underlying model if required
if serialize_model_using_mlflow:
underlying_model_flavor = get_underlying_model_flavor(explainer.model)
if underlying_model_flavor != _UNKNOWN_MODEL_FLAVOR:
serializable_by_mlflow = True # prevents SHAP from serializing the underlying model
underlying_model_path = os.path.join(path,
_UNDERLYING_MODEL_SUBPATH)
else:
warnings.warn(
"Unable to serialize underlying model using MLflow, will use SHAP serialization"
)
if underlying_model_flavor == mlflow.sklearn.FLAVOR_NAME:
mlflow.sklearn.save_model(explainer.model.inner_model.__self__,
underlying_model_path)
elif underlying_model_flavor == mlflow.pytorch.FLAVOR_NAME:
mlflow.pytorch.save_model(explainer.model.inner_model,
underlying_model_path)
# saving the explainer object
explainer_data_subpath = "explainer.shap"
explainer_output_path = os.path.join(path, explainer_data_subpath)
with open(explainer_output_path, "wb") as explainer_output_file_handle:
if serialize_model_using_mlflow and serializable_by_mlflow:
explainer.save(explainer_output_file_handle, model_saver=False)
else:
explainer.save(explainer_output_file_handle)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.shap",
model_path=explainer_data_subpath,
underlying_model_flavor=underlying_model_flavor,
env=_CONDA_ENV_FILE_NAME,
code=code_dir_subpath,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
shap_version=shap.__version__,
serialized_explainer=explainer_data_subpath,
underlying_model_flavor=underlying_model_flavor,
code=code_dir_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
if underlying_model_path is not None:
underlying_model_conda_env = _get_conda_env_for_underlying_model(
underlying_model_path)
conda_env = _merge_environments(conda_env, underlying_model_conda_env)
pip_requirements = _get_pip_deps(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
_PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))
def save_model(
sk_model,
path,
conda_env=None,
mlflow_model=None,
serialization_format=SERIALIZATION_FORMAT_CLOUDPICKLE,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
"""
Save a scikit-learn model to a path on the local file system. Produces an MLflow Model
containing the following flavors:
- :py:mod:`mlflow.sklearn`
- :py:mod:`mlflow.pyfunc`. NOTE: This flavor is only included for scikit-learn models
that define `predict()`, since `predict()` is required for pyfunc model inference.
:param sk_model: scikit-learn model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is 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',
'scikit-learn=0.19.2'
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param serialization_format: The format in which to serialize the model. This should be one of
the formats listed in
``mlflow.sklearn.SUPPORTED_SERIALIZATION_FORMATS``. The Cloudpickle
format, ``mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE``,
provides better cross-system compatibility by identifying and
packaging code dependencies with the serialized model.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
.. code-block:: python
:caption: Example
import mlflow.sklearn
from sklearn.datasets import load_iris
from sklearn import tree
iris = load_iris()
sk_model = tree.DecisionTreeClassifier()
sk_model = sk_model.fit(iris.data, iris.target)
# Save the model in cloudpickle format
# set path to location for persistence
sk_path_dir_1 = ...
mlflow.sklearn.save_model(
sk_model, sk_path_dir_1,
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE)
# save the model in pickle format
# set path to location for persistence
sk_path_dir_2 = ...
mlflow.sklearn.save_model(sk_model, sk_path_dir_2,
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE)
"""
import sklearn
if serialization_format not in SUPPORTED_SERIALIZATION_FORMATS:
raise MlflowException(
message=
("Unrecognized serialization format: {serialization_format}. Please specify one"
" of the following supported formats: {supported_formats}.".
format(
serialization_format=serialization_format,
supported_formats=SUPPORTED_SERIALIZATION_FORMATS,
)),
error_code=INVALID_PARAMETER_VALUE,
)
if os.path.exists(path):
raise MlflowException(message="Path '{}' already exists".format(path),
error_code=RESOURCE_ALREADY_EXISTS)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "model.pkl"
_save_model(
sk_model=sk_model,
output_path=os.path.join(path, model_data_subpath),
serialization_format=serialization_format,
)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env(
include_cloudpickle=serialization_format ==
SERIALIZATION_FORMAT_CLOUDPICKLE)
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# `PyFuncModel` only works for sklearn models that define `predict()`.
if hasattr(sk_model, "predict"):
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.sklearn",
model_path=model_data_subpath,
env=conda_env_subpath,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
pickled_model=model_data_subpath,
sklearn_version=sklearn.__version__,
serialization_format=serialization_format,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(
statsmodels_model,
path,
conda_env=None,
mlflow_model=None,
remove_data: bool = False,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
"""
Save a statsmodels model to a path on the local file system.
:param statsmodels_model: statsmodels model (an instance of `statsmodels.base.model.Results`_)
to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this describes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is 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',
'statsmodels=0.11.1'
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param remove_data: bool. If False (default), then the instance is pickled without changes.
If True, then all arrays with length nobs are set to None before
pickling. See the remove_data method.
In some cases not all arrays will be set to None.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
"""
import statsmodels
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
model_data_path = os.path.join(path, STATSMODELS_DATA_SUBPATH)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save a statsmodels model
statsmodels_model.save(model_data_path, remove_data)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.statsmodels",
data=STATSMODELS_DATA_SUBPATH,
env=conda_env_subpath,
)
mlflow_model.add_flavor(FLAVOR_NAME,
statsmodels_version=statsmodels.__version__,
data=STATSMODELS_DATA_SUBPATH)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
