mlflow.pytorch

The mlflow.pytorch 模块提供用于记录和加载 PyTorch 模型的 API。该模块 以以下 flavors 导出 PyTorch 模型：

PyTorch (native) format

这是可以加载回 PyTorch 的主要 flavor。

mlflow.pyfunc

用于通用的基于 pyfunc 的部署工具和批量推理。

class mlflow.pytorch.MlflowModelCheckpointCallback(monitor='val_loss', mode='min', save_best_only=True, save_weights_only=False, save_freq='epoch')[source]

基类： pytorch_lightning.callbacks.callback.Callback, mlflow.utils.checkpoint_utils.MlflowModelCheckpointCallbackBase

用于将 pytorch-lightning 的模型检查点自动记录到 MLflow 的回调。 此回调实现仅支持 pytorch-lightning >= 1.6.0。

Parameters

  • monitor – 在自动模型检查点保存中，如果将 model_checkpoint_save_best_only 设置为 True，则要监控的指标名称。

  • save_best_only – 如果 True，则只有当模型根据被监控的量被认为是“最佳”模型时，自动模型检查点才会保存，并且之前的检查点模型将被覆盖。

  • mode – 值为 {“min”, “max”} 之一。 在自动模型检查点保存中，如果 save_best_only=True，是否覆盖当前保存文件的决定是基于对被监测指标的最大化或最小化。

  • save_weights_only – 在自动模型检查点保存中，如果 True，则只会保存模型的权重。否则，优化器状态、lr-scheduler 状态等也会被加入到检查点中。

  • save_freq – “epoch” 或整数。当使用 “epoch” 时，回调在每个 epoch 之后保存模型。当使用整数时，回调在经过这么多批次后保存模型。注意，如果保存不是与 epoch 对齐，被监控的指标可能不那么可靠（它可能只反映最少 1 个批次，因为指标会在每个 epoch 重置）。默认值为 “epoch”。

Example 
import mlflow
from mlflow.pytorch import MlflowModelCheckpointCallback
from pytorch_lightning import Trainer

mlflow.pytorch.autolog(checkpoint=True)

model = MyLightningModuleNet()  # A custom-pytorch lightning model
train_loader = create_train_dataset_loader()

mlflow_checkpoint_callback = MlflowModelCheckpointCallback()

trainer = Trainer(callbacks=[mlflow_checkpoint_callback])

with mlflow.start_run() as run:
    trainer.fit(model, train_loader)
on_fit_start(trainer: pytorch_lightning.trainer.trainer.Trainer, pl_module: pytorch_lightning.core.module.LightningModule) → None[source]

在 fit 开始时调用。

on_train_batch_end(trainer: pytorch_lightning.trainer.trainer.Trainer, pl_module: pytorch_lightning.core.module.LightningModule, outputs, batch, batch_idx) → None[source]

在训练批次结束时被调用。

注意

这里的 outputs["loss"] 的值将是相对于 accumulate_grad_batches 对从 training_step 返回的损失进行归一化后的值。

on_train_epoch_end(trainer: pytorch_lightning.trainer.trainer.Trainer, pl_module: pytorch_lightning.core.module.LightningModule) → None[source]

在训练 epoch 结束时调用。

要在 epoch 结束时访问所有批次的输出，你可以将 step 输出缓存为 pytorch_lightning.core.LightningModule 的一个属性，并在此钩子中访问它：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss


class MyCallback(L.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
        pl_module.log("training_epoch_mean", epoch_mean)
        # free up the memory
        pl_module.training_step_outputs.clear()
save_checkpoint(filepath: str)[source]
mlflow.pytorch.autolog(log_every_n_epoch=1, log_every_n_step=None, log_models=True, log_datasets=True, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False, registered_model_name=None, extra_tags=None, checkpoint=True, checkpoint_monitor='val_loss', checkpoint_mode='min', checkpoint_save_best_only=True, checkpoint_save_weights_only=False, checkpoint_save_freq='epoch')[source]

注意

Autologging 已知与以下包版本兼容：2.1.0 <= torch <= 2.8.0。在此范围之外的包版本可能导致 Autologging 无法成功。

启用（或禁用）并配置从 PyTorch Lightning 到 MLflow 的自动记录。

当您调用 pytorch_lightning.Trainer() 的 fit 方法时，会执行自动记录。

探索完整的 PyTorch MNIST 以获取一个包含额外 lightening 步骤实现的详尽示例。

注意：完整的自动记录仅支持 PyTorch Lightning 模型，即继承自 pytorch_lightning.LightningModule 的模型。对于原生 PyTorch（即仅继承自 torch.nn.Module 的模型），自动记录仅会将对 torch.utils.tensorboard.SummaryWriter 的 add_scalar 和 add_hparams 方法的调用记录到 mlflow。在这种情况下，也没有“epoch”的概念。

Parameters

  • log_every_n_epoch – 如果指定，则每隔 n 个 epoch 记录一次指标。默认情况下，指标在每个 epoch 之后记录。

  • log_every_n_step – 如果指定，则每隔 n 个训练步骤记录一次批次指标。默认情况下，不会对训练步骤记录指标。请注意，将其设置为 1 可能会导致性能问题，不建议这样做。指标是针对 Lightning 的全局步骤编号记录的；当使用多个优化器时，假定在每个训练步骤中所有优化器都会被更新。

  • log_models – 如果 True，已训练的模型会被记录为 MLflow 模型工件。 如果 False，已训练的模型不会被记录。

  • log_datasets – 如果 True，则将数据集信息记录到 MLflow Tracking。 如果 False，则不记录数据集信息。

  • disable – 如果 True，则禁用 PyTorch Lightning 的 autologging 集成。 如果 False，则启用 PyTorch Lightning 的 autologging 集成。

  • exclusive – 如果 True，自动记录的内容不会记录到用户创建的 fluent 运行中。 如果 False，自动记录的内容会记录到活动的 fluent 运行，该运行可能是用户创建的。

  • disable_for_unsupported_versions – 如果 True，则禁用对未与此版本的 MLflow 客户端测试过或与之不兼容的 pytorch 和 pytorch-lightning 版本的自动记录（autologging）。

  • silent – 如果 True，在 PyTorch Lightning 自动记录（autologging）期间抑制来自 MLflow 的所有事件日志和警告。如果 False，在 PyTorch Lightning 自动记录期间显示所有事件和警告。

  • registered_model_name – 如果提供，每次训练模型时，该模型都会作为具有此名称的注册模型的新模型版本进行注册。如果该注册模型尚不存在，则会创建它。

  • extra_tags – 一个字典，用于在 autologging 创建的每个托管运行上设置额外标签。

  • checkpoint – 启用自动模型检查点，该功能仅支持 pytorch-lightning >= 1.6.0。

  • checkpoint_monitor – 在自动模型检查点保存中，如果将 model_checkpoint_save_best_only 设置为 True，则要监控的指标名称。

  • checkpoint_mode – one of {"min", "max"}. 在自动模型检查点保存中， 如果 save_best_only=True，是否覆盖当前保存文件的决策基于对被监测指标的最大化或最小化。

  • checkpoint_save_best_only – If True, 自动模型检查点保存只有在根据所监控的量认为模型是“最佳”模型时才会进行，并且会覆盖之前的检查点模型。

  • checkpoint_save_weights_only – 在自动模型检查点保存中，如果 True，则仅保存模型的权重。否则，优化器状态、学习率调度器状态等也会被添加到检查点中。

  • checkpoint_save_freq – “epoch” 或整数。 当使用 “epoch” 时，回调会在每个 epoch 之后保存模型。 当使用整数时，回调将在这么多批次结束时保存模型。 请注意，如果保存点与 epochs 未对齐，被监控的指标可能不太可靠（它可能只反映最少 1 个批次，因为指标会在每个 epoch 重置）。 默认值为 “epoch”。

Example 
import os

import lightning as L
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader, Subset
from torchmetrics import Accuracy
from torchvision import transforms
from torchvision.datasets import MNIST

import mlflow.pytorch
from mlflow import MlflowClient


class MNISTModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)
        self.accuracy = Accuracy("multiclass", num_classes=10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        pred = logits.argmax(dim=1)
        acc = self.accuracy(pred, y)

        # PyTorch `self.log` will be automatically captured by MLflow.
        self.log("train_loss", loss, on_epoch=True)
        self.log("acc", acc, on_epoch=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


def print_auto_logged_info(r):
    tags = {k: v for k, v in r.data.tags.items() if not k.startswith("mlflow.")}
    artifacts = [f.path for f in MlflowClient().list_artifacts(r.info.run_id, "model")]
    print(f"run_id: {r.info.run_id}")
    print(f"artifacts: {artifacts}")
    print(f"params: {r.data.params}")
    print(f"metrics: {r.data.metrics}")
    print(f"tags: {tags}")


# Initialize our model.
mnist_model = MNISTModel()

# Load MNIST dataset.
train_ds = MNIST(
    os.getcwd(), train=True, download=True, transform=transforms.ToTensor()
)
# Only take a subset of the data for faster training.
indices = torch.arange(32)
train_ds = Subset(train_ds, indices)
train_loader = DataLoader(train_ds, batch_size=8)

# Initialize a trainer.
trainer = L.Trainer(max_epochs=3)

# Auto log all MLflow entities
mlflow.pytorch.autolog()

# Train the model.
with mlflow.start_run() as run:
    trainer.fit(mnist_model, train_loader)

# Fetch the auto logged parameters and metrics.
print_auto_logged_info(mlflow.get_run(run_id=run.info.run_id))
mlflow.pytorch.get_default_conda_env()[source]
Returns

作为字典形式的默认 Conda 环境，用于通过调用 save_model() 和 log_model() 生成的 MLflow Models。

Example 
import mlflow

# Log PyTorch model
with mlflow.start_run() as run:
    mlflow.pytorch.log_model(model, name="model", signature=signature)

# Fetch the associated conda environment
env = mlflow.pytorch.get_default_conda_env()
print(f"conda env: {env}")
Output 
conda env {'name': 'mlflow-env',
           'channels': ['conda-forge'],
           'dependencies': ['python=3.8.15',
                            {'pip': ['torch==1.5.1',
                                     'mlflow',
                                     'cloudpickle==1.6.0']}]}
mlflow.pytorch.get_default_pip_requirements()[source]
Returns

此 flavor 生成的 MLflow Models 的默认 pip 依赖项列表。调用 save_model() 和 log_model() 会生成一个 pip 环境，该环境至少包含这些依赖项。

mlflow.pytorch.load_checkpoint(model_class, run_id=None, epoch=None, global_step=None, kwargs=None)[source]

如果在 autologging 中启用“checkpoint”，在 pytorch-lightning 模型训练执行期间，被 checkpoint 的模型会作为 MLflow 工件记录。使用此 API，您可以加载被 checkpoint 的模型。

如果你想加载最新的检查点，请将 epoch 和 global_step 都设置为 None。 如果在 autologging 中将 “checkpoint_save_freq” 设置为 “epoch”，你可以将 epoch 参数设置为要加载的检查点的 epoch，以加载特定 epoch 的检查点。 如果在 autologging 中将 “checkpoint_save_freq” 设置为一个整数，你可以将 global_step 参数设置为要加载的检查点的 global step，以加载特定 global step 的检查点。 epoch 参数和 global_step 不能同时设置。

Parameters

  • model_class – 训练模型的类，类应该继承‘pytorch_lightning.LightningModule’。

  • run_id – 模型被记录到的运行的 id。如果未提供，则使用当前活动的运行。

  • epoch – 要加载的检查点的 epoch，如果您将 “checkpoint_save_freq” 设置为 “epoch”。

  • global_step – 要加载的检查点的全局步骤，如果你将 “checkpoint_save_freq” 设置为整数。

  • kwargs – 初始化模型所需的任何额外 kwargs。

Returns

从指定的检查点恢复的 pytorch-lightning 模型的实例。

Example 
import mlflow

mlflow.pytorch.autolog(checkpoint=True)

model = MyLightningModuleNet()  # A custom-pytorch lightning model
train_loader = create_train_dataset_loader()
trainer = Trainer()

with mlflow.start_run() as run:
    trainer.fit(model, train_loader)

run_id = run.info.run_id

# load latest checkpoint model
latest_checkpoint_model = mlflow.pytorch.load_checkpoint(MyLightningModuleNet, run_id)

# load history checkpoint model logged in second epoch
checkpoint_model = mlflow.pytorch.load_checkpoint(MyLightningModuleNet, run_id, epoch=2)
mlflow.pytorch.load_model(model_uri, dst_path=None, **kwargs)[source]

从本地文件或某次运行中加载 PyTorch 模型。

Parameters

  • model_uri –

    MLflow 模型的 URI 格式的位置，例如：

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs://run-relative/path/to/model

    • models://

    • models://

    有关受支持的 URI 方案的更多信息，请参见 Referencing Artifacts。

  • dst_path – 下载模型工件的本地文件系统路径。该目录必须已存在。如果未指定，将创建一个本地输出路径。

  • kwargs – 要传递给 torch.load 方法的 kwargs。

Returns

一个 PyTorch 模型。

Example 
import torch
import mlflow.pytorch


model = nn.Linear(1, 1)

# Log the model
with mlflow.start_run() as run:
    mlflow.pytorch.log_model(model, name="model")

# Inference after loading the logged model
model_uri = f"runs:/{run.info.run_id}/model"
loaded_model = mlflow.pytorch.load_model(model_uri)
for x in [4.0, 6.0, 30.0]:
    X = torch.Tensor([[x]])
    y_pred = loaded_model(X)
    print(f"predict X: {x}, y_pred: {y_pred.data.item():.2f}")
Output 
predict X: 4.0, y_pred: 7.57
predict X: 6.0, y_pred: 11.64
predict X: 30.0, y_pred: 60.48
mlflow.pytorch.log_model(pytorch_model, artifact_path: str | None = None, conda_env=None, code_paths=None, pickle_module=None, registered_model_name=None, signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, await_registration_for=300, extra_files=None, pip_requirements=None, extra_pip_requirements=None, metadata=None, name: str | None = None, params: dict[str, typing.Any] | None = None, tags: dict[str, typing.Any] | None = None, model_type: str | None = None, step: int = 0, model_id: str | None = None, **kwargs)[source]

将 PyTorch 模型记录为当前运行的 MLflow 工件。

警告

使用签名记录模型以避免推理错误。如果记录模型时没有提供签名，MLflow Model Server 会依赖 NumPy 推断出的默认数据类型。然而，PyTorch 通常期望不同的默认值，尤其是在解析浮点数时。必须包含签名，以确保模型以正确的数据类型被记录，从而使 MLflow model server 能够正确提供有效输入。

Parameters

  • pytorch_model –

    要保存的 PyTorch 模型。可以是 eager 模型（torch.nn.Module 的子类），也可以是通过 torch.jit.script 或 torch.jit.trace 准备的脚本化模型。

    该模型接受单个 torch.FloatTensor 作为输入并产生单个输出张量。

    如果保存的是 eager 模型，模型类的任何代码依赖（包括类定义本身）应包含在以下位置之一：

    • 模型 Conda 环境中列出的包，由 conda_env 参数指定。

    • 由 code_paths 参数指定的一个或多个文件。

  • artifact_path – 已弃用。请改用 name。

  • conda_env –

    Conda 环境的字典表示或 conda environment yaml 文件的路径。如果提供，描述了运行该模型应当使用的环境。至少，它应当指定 get_default_conda_env() 中包含的依赖项。如果 None，则会向模型添加一个 conda 环境，其 pip 依赖项由 mlflow.models.infer_pip_requirements() 推断得到。如果依赖推断失败，则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖项会被写入到一个 pip requirements.txt 文件中，完整的 conda 环境会写入到 conda.yaml。 下面是一个 示例 的 conda 环境字典表示：

    {
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "torch==x.y.z"
            ],
        },
    ],
}

  • code_paths –

    本地文件系统中指向 Python 文件依赖（或包含文件依赖的目录）路径的列表。这些文件在模型加载时会被预先添加到系统路径中。如果为某个模型声明了依赖文件且多个文件之间存在导入依赖关系，那么这些文件应从一个共同的根路径声明相对导入，以避免在加载模型时发生导入错误。

    有关 code_paths 功能、推荐的使用模式和限制的详细说明，请参阅 code_paths usage guide。

  • pickle_module – PyTorch 用于序列化（“pickle”）指定的 pytorch_model 的模块。该模块作为 pickle_module 参数传递给 torch.save()。默认情况下，该模块在加载时也用于反序列化（“unpickle”）PyTorch 模型。

  • registered_model_name – 如果提供，则在 registered_model_name 下创建模型版本，如果不存在具有给定名称的注册模型，也会创建一个注册模型。

  • signature –

    一个 ModelSignature 类的实例，用于描述模型的输入和输出。如果未指定但提供了 input_example，则会根据提供的输入示例和模型自动推断签名。要在提供输入示例时禁用自动签名推断，请将 signature 设置为 False。要手动推断模型签名，请在具有有效模型输入，例如省略了目标列的训练数据集，以及有效模型输出，例如在训练数据集上生成的模型预测的数据集上调用 infer_signature()，例如：

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # 计算模型预测
signature = infer_signature(train, predictions)

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame，然后使用 Pandas 的 split-oriented 格式序列化为 json，或者转换为一个 numpy array，其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • await_registration_for – 等待模型版本创建完成并处于 READY 状态的秒数。默认情况下，该函数等待五分钟。指定 0 或 None 可跳过等待。

  • extra_files –

    一个包含对应额外文件路径的列表，如果 None，则不会向模型添加额外文件。远程 URI 会解析为绝对文件系统路径。例如，考虑以下 extra_files 列表：

    extra_files = ["s3://my-bucket/path/to/my_file1", "s3://my-bucket/path/to/my_file2"]

    在这种情况下，来自 S3 的 "my_file1 & my_file2" 额外文件会被下载。

  • pip_requirements – 要么是 pip 需求字符串的可迭代对象（例如 ["torch", "-r requirements.txt", "-c constraints.txt"]），要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，它描述了该模型应该运行的环境。如果 None，则会由 mlflow.models.infer_pip_requirements() 从当前软件环境推断出默认的依赖列表。如果依赖推断失败，则回退使用 get_default_pip_requirements。依赖和约束会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分进行存储。依赖还会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

  • extra_pip_requirements –

    要么是一个 pip 需求字符串的可迭代对象（例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]），要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该参数描述了附加的 pip 依赖，这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定：

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip 依赖。

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • name – 模型名称。

  • params – 一个用于与模型一同记录的参数字典。

  • tags – 一个要与模型一起记录的标签字典。

  • model_type – 模型的类型。

  • step – 在该步记录模型输出和指标

  • model_id – 模型的 ID。

  • kwargs – 传递给 torch.save 方法的 kwargs。

Returns

一个 ModelInfo 实例，包含已记录模型的元数据。

Example 
import numpy as np
import torch
import mlflow
from mlflow import MlflowClient
from mlflow.models import infer_signature

# Define model, loss, and optimizer
model = nn.Linear(1, 1)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# Create training data with relationship y = 2X
X = torch.arange(1.0, 26.0).reshape(-1, 1)
y = X * 2

# Training loop
epochs = 250
for epoch in range(epochs):
    # Forward pass: Compute predicted y by passing X to the model
    y_pred = model(X)

    # Compute the loss
    loss = criterion(y_pred, y)

    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Create model signature
signature = infer_signature(X.numpy(), model(X).detach().numpy())

# Log the model
with mlflow.start_run() as run:
    mlflow.pytorch.log_model(model, name="model")

    # convert to scripted model and log the model
    scripted_pytorch_model = torch.jit.script(model)
    mlflow.pytorch.log_model(scripted_pytorch_model, name="scripted_model")

# Fetch the logged model artifacts
print(f"run_id: {run.info.run_id}")
for artifact_path in ["model/data", "scripted_model/data"]:
    artifacts = [
        f.path for f in MlflowClient().list_artifacts(run.info.run_id, artifact_path)
    ]
    print(f"artifacts: {artifacts}")
Output 
run_id: 1a1ec9e413ce48e9abf9aec20efd6f71
artifacts: ['model/data/model.pth',
            'model/data/pickle_module_info.txt']
artifacts: ['scripted_model/data/model.pth',
            'scripted_model/data/pickle_module_info.txt']
../_images/pytorch_logged_models.png

PyTorch 已记录的模型

mlflow.pytorch.save_model(pytorch_model, path, conda_env=None, mlflow_model=None, code_paths=None, pickle_module=None, signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, extra_files=None, pip_requirements=None, extra_pip_requirements=None, metadata=None, **kwargs)[source]

将 PyTorch 模型保存到本地文件系统上的某个路径。

Parameters

  • pytorch_model –

    要保存的 PyTorch 模型。可以是一个 eager 模型（torch.nn.Module 的子类），也可以是通过 torch.jit.script 或 torch.jit.trace 准备的脚本化模型。

    要保存 eager 模型，模型类的任何代码依赖项（包括类定义本身）应包含在以下任一位置：

    • 模型 Conda 环境中列出的包，由 conda_env 参数指定。

    • 由 code_paths 参数指定的一个或多个文件。

  • path – 模型要保存的本地路径。

  • conda_env –

    要么是 Conda 环境的字典表示，要么是指向 conda 环境 yaml 文件的路径。如果提供，则描述该模型应运行的环境。至少，它应当指定 get_default_conda_env() 中包含的依赖项。如果 None，则会向模型添加一个 conda 环境，该环境的 pip 依赖由 mlflow.models.infer_pip_requirements() 推断。如果依赖推断失败，则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖会被写入到 pip 文件 requirements.txt，完整的 conda 环境会被写入到 conda.yaml。以下是 conda 环境的一个 示例 字典表示：

    {
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "torch==x.y.z"
            ],
        },
    ],
}

  • mlflow_model – mlflow.models.Model 此 flavor 正在被添加到。

  • code_paths –

    本地文件系统中指向 Python 文件依赖（或包含文件依赖的目录）路径的列表。这些文件在模型加载时会被预先添加到系统路径中。如果为某个模型声明了依赖文件且多个文件之间存在导入依赖关系，那么这些文件应从一个共同的根路径声明相对导入，以避免在加载模型时发生导入错误。

    有关 code_paths 功能、推荐的使用模式和限制的详细说明，请参阅 code_paths usage guide。

  • pickle_module – PyTorch 应使用来序列化（“pickle”）指定的 pytorch_model 的模块。该模块作为 pickle_module 参数传递给 torch.save()。默认情况下，该模块也用于在加载时反序列化（“unpickle”）模型。

  • signature –

    一个 ModelSignature 类的实例，用于描述模型的输入和输出。如果未指定但提供了 input_example，则会根据提供的输入示例和模型自动推断签名。要在提供输入示例时禁用自动签名推断，请将 signature 设置为 False。要手动推断模型签名，请在具有有效模型输入，例如省略了目标列的训练数据集，以及有效模型输出，例如在训练数据集上生成的模型预测的数据集上调用 infer_signature()，例如：

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # 计算模型预测
signature = infer_signature(train, predictions)

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame，然后使用 Pandas 的 split-oriented 格式序列化为 json，或者转换为一个 numpy array，其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • extra_files –

    包含对应额外文件路径的列表。远程 URI 会被解析为绝对文件系统路径。 例如，考虑下列 extra_files 列表 -

    extra_files = [“s3://my-bucket/path/to/my_file1”, “s3://my-bucket/path/to/my_file2”]

    在这种情况下，"my_file1 & my_file2" 额外文件会从 S3 下载。

    如果 None，则不会向模型添加额外文件。

  • pip_requirements – 要么是可迭代的 pip 依赖字符串（例如 ["torch", "-r requirements.txt", "-c constraints.txt"]）要么是本地文件系统中 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该项描述了此模型应运行的环境。如果 None，则会由 mlflow.models.infer_pip_requirements() 根据当前软件环境推断出默认的依赖列表。如果依赖推断失败，则回退到使用 get_default_pip_requirements。依赖和约束会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

  • extra_pip_requirements –

    要么是一个 pip 需求字符串的可迭代对象（例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]），要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该参数描述了附加的 pip 依赖，这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定：

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip 依赖。

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • kwargs – 要传递给 torch.save 方法的 kwargs。

Example 
import os
import mlflow
import torch


model = nn.Linear(1, 1)

# 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 = f"{os.getcwd()}/{model_path}"
    loaded_model = mlflow.pytorch.load_model(model_uri)
    print(f"Loaded {model_path}:")
    for x in [6.0, 8.0, 12.0, 30.0]:
        X = torch.Tensor([[x]])
        y_pred = loaded_model(X)
        print(f"predict X: {x}, y_pred: {y_pred.data.item():.2f}")
    print("--")
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