MLflow PyTorch 集成

简介

PyTorch 是由 Meta AI 研究实验室开发的开源深度学习框架。它提供了动态计算图和用于构建神经网络的 Python 风格 API，使其在研究和生产深度学习应用中广受欢迎。

MLflow 的 PyTorch 集成功能为深度学习工作流提供了实验跟踪、模型版本控制和部署能力。

为什么选择 MLflow + PyTorch？

自动日志记录

只需一行代码即可实现全面的实验跟踪：mlflow.pytorch.autolog() 会自动记录指标、参数和模型。

实验跟踪

跟踪所有 PyTorch 实验中的训练指标、超参数、模型架构和工件（artifacts）。

模型注册表

利用 MLflow 的模型注册表（Model Registry）和服务基础设施，对 PyTorch 模型进行版本控制、分阶段管理和部署。

可复现性

捕获模型状态、随机种子和环境，以实现可复现的深度学习实验。

入门

只需一行代码即可实现全面的实验跟踪

python

import mlflow
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Enable autologging
mlflow.pytorch.autolog()

# Create synthetic data
X = torch.randn(1000, 784)
y = torch.randint(0, 10, (1000,))
train_loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)

# Your existing PyTorch code works unchanged
model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Training loop - metrics, parameters, and models logged automatically
for epoch in range(10):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

自动记录（Autologging）功能会自动捕获训练指标、模型参数、优化器配置和模型检查点。

支持 PyTorch Lightning

MLflow 的自动记录功能可与 PyTorch Lightning 无缝配合。对于使用自定义训练循环的原生 PyTorch，请使用如下文所示的手动记录方式。

手动记录

对于标准 PyTorch 工作流，将 MLflow 记录功能集成到您的训练循环中

python

import mlflow
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader


# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        return self.linear_relu_stack(x)


# Training parameters
params = {
    "epochs": 5,
    "learning_rate": 1e-3,
    "batch_size": 64,
}

# Training with MLflow logging
with mlflow.start_run():
    # Log parameters
    mlflow.log_params(params)

    # Initialize model and optimizer
    model = NeuralNetwork()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=params["learning_rate"])

    # Training loop
    for epoch in range(params["epochs"]):
        model.train()
        train_loss = 0
        correct = 0
        total = 0

        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = loss_fn(output, target)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()

        # Log metrics per epoch
        avg_loss = train_loss / len(train_loader)
        accuracy = 100.0 * correct / total

        mlflow.log_metrics({"train_loss": avg_loss, "train_accuracy": accuracy}, step=epoch)

    # Log final model
    mlflow.pytorch.log_model(model, name="model")

系统指标跟踪

在训练期间跟踪硬件资源利用率，以监控 GPU 使用情况、内存消耗和系统性能

python

import mlflow
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Create data and model
X = torch.randn(1000, 784)
y = torch.randint(0, 10, (1000,))
train_loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)

model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Enable system metrics logging
mlflow.enable_system_metrics_logging()

with mlflow.start_run():
    mlflow.log_params({"learning_rate": 0.001, "batch_size": 32, "epochs": 10})

    # Training loop - system metrics logged automatically
    for epoch in range(10):
        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

        mlflow.log_metric("loss", loss.item(), step=epoch)

    mlflow.pytorch.log_model(model, name="model")

系统指标记录会自动捕获

• GPU 指标：利用率百分比、内存使用量、温度和功耗
• CPU 指标：利用率百分比和内存使用量
• 磁盘 I/O：读/写吞吐量和利用率
• 网络 I/O：网络流量统计

对于 PyTorch 训练，监控 GPU 指标对于以下方面特别有价值：

• 识别可能暗示数据加载瓶颈的 GPU 利用率不足问题
• 在发生内存不足（OOM）错误之前检测内存问题
• 根据 GPU 内存使用情况优化批处理大小
• 比较不同模型架构间的资源效率

高级配置

您可以自定义系统指标的采集频率和行为。有关详细的配置选项，请参阅系统指标文档。

带签名的模型记录

使用输入/输出签名记录 PyTorch 模型，以便更好地理解模型

python

import mlflow
import torch
import torch.nn as nn
from mlflow.models import infer_signature

model = nn.Sequential(nn.Linear(10, 50), nn.ReLU(), nn.Linear(50, 1))

# Create sample input and output for signature
input_example = torch.randn(1, 10)
predictions = model(input_example)

# Infer signature from input/output
signature = infer_signature(input_example.numpy(), predictions.detach().numpy())

with mlflow.start_run():
    # Log model with signature and input example
    mlflow.pytorch.log_model(
        model,
        name="pytorch_model",
        signature=signature,
        input_example=input_example.numpy(),
    )

检查点跟踪

使用 MLflow 3 的检查点版本控制在训练期间跟踪模型检查点。使用 step 参数对检查点进行版本控制，并将指标链接到特定的模型版本

python

import mlflow
import torch
import torch.nn as nn
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split


# Helper function to prepare data
def prepare_data(df):
    X = torch.tensor(df.iloc[:, :-1].values, dtype=torch.float32)
    y = torch.tensor(df.iloc[:, -1].values, dtype=torch.long)
    return X, y


# Helper function to compute accuracy
def compute_accuracy(model, X, y):
    with torch.no_grad():
        outputs = model(X)
        _, predicted = torch.max(outputs, 1)
        accuracy = (predicted == y).sum().item() / y.size(0)
    return accuracy


# Define a basic PyTorch classifier
class IrisClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x


# Load Iris dataset and prepare the DataFrame
iris = load_iris()
iris_df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
iris_df["target"] = iris.target

# Split into training and testing datasets
train_df, test_df = train_test_split(iris_df, test_size=0.2, random_state=42)

# Prepare training data
train_dataset = mlflow.data.from_pandas(train_df, name="iris_train")
X_train, y_train = prepare_data(train_dataset.df)

# Initialize model
input_size = X_train.shape[1]
hidden_size = 16
output_size = len(iris.target_names)
model = IrisClassifier(input_size, hidden_size, output_size)

# Training configuration
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

with mlflow.start_run() as run:
    # Log parameters once at the start
    mlflow.log_params({
        "n_layers": 3,
        "activation": "ReLU",
        "criterion": "CrossEntropyLoss",
        "optimizer": "Adam",
        "learning_rate": 0.01,
    })

    for epoch in range(101):
        # Training step
        out = model(X_train)
        loss = criterion(out, y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Log a checkpoint every 10 epochs
        if epoch % 10 == 0:
            # Log model checkpoint with step parameter
            model_info = mlflow.pytorch.log_model(
                pytorch_model=model,
                name=f"iris-checkpoint-{epoch}",
                step=epoch,
                input_example=X_train[:5].numpy(),
            )

            # Log metrics linked to this checkpoint and dataset
            accuracy = compute_accuracy(model, X_train, y_train)
            mlflow.log_metric(
                key="train_accuracy",
                value=accuracy,
                step=epoch,
                model_id=model_info.model_id,
                dataset=train_dataset,
            )

# Search and rank checkpoints by performance
ranked_checkpoints = mlflow.search_logged_models(
    filter_string=f"source_run_id='{run.info.run_id}'",
    order_by=[{"field_name": "metrics.train_accuracy", "ascending": False}],
    output_format="list",
)

best_checkpoint = ranked_checkpoints[0]
print(f"Best checkpoint: {best_checkpoint.name}")
print(f"Accuracy: {best_checkpoint.metrics[0].value}")

此方法使您能够：

• 系统化地对检查点进行版本控制（使用 step 参数）
• 将指标链接到特定模型（在 log_metric() 中使用 model_id）
• 将指标与数据集关联，以实现更好的跟踪
• 对检查点进行排名和比较（使用 mlflow.search_logged_models()）

模型加载

加载已记录的 PyTorch 模型以进行推理

python

# Load as PyTorch model
model_uri = "runs:/<run_id>/pytorch_model"
loaded_model = mlflow.pytorch.load_model(model_uri)

# Make predictions
input_tensor = torch.randn(5, 10)
predictions = loaded_model(input_tensor)

# Load as PyFunc for generic inference
pyfunc_model = mlflow.pyfunc.load_model(model_uri)
predictions = pyfunc_model.predict(input_tensor.numpy())

超参数优化

使用 MLflow 跟踪超参数调整实验

python

import mlflow
import optuna
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Create synthetic dataset for demonstration
input_size = 784  # e.g., flattened 28x28 images
output_size = 10  # e.g., 10 classes

X_train = torch.randn(1000, input_size)
y_train = torch.randint(0, output_size, (1000,))
X_val = torch.randn(200, input_size)
y_val = torch.randint(0, output_size, (200,))

train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32)


def train_and_evaluate(model, optimizer, train_loader, val_loader, epochs=5):
    """Simple training loop for demonstration."""
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

    # Validation
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for data, target in val_loader:
            output = model(data)
            val_loss += criterion(output, target).item()

    return val_loss / len(val_loader)


def objective(trial):
    """Optuna objective for hyperparameter tuning."""

    with mlflow.start_run(nested=True):
        # Define hyperparameter search space
        params = {
            "learning_rate": trial.suggest_float("learning_rate", 1e-5, 1e-1, log=True),
            "hidden_size": trial.suggest_int("hidden_size", 32, 512),
            "dropout": trial.suggest_float("dropout", 0.1, 0.5),
        }

        # Log parameters
        mlflow.log_params(params)

        # Create model
        model = nn.Sequential(
            nn.Linear(input_size, params["hidden_size"]),
            nn.ReLU(),
            nn.Dropout(params["dropout"]),
            nn.Linear(params["hidden_size"], output_size),
        )

        # Train model
        optimizer = optim.Adam(model.parameters(), lr=params["learning_rate"])
        val_loss = train_and_evaluate(model, optimizer, train_loader, val_loader)

        # Log validation loss
        mlflow.log_metric("val_loss", val_loss)

        return val_loss


# Run optimization
with mlflow.start_run(run_name="PyTorch HPO"):
    study = optuna.create_study(direction="minimize")
    study.optimize(objective, n_trials=50)

    # Log best parameters
    mlflow.log_params({f"best_{k}": v for k, v in study.best_params.items()})
    mlflow.log_metric("best_val_loss", study.best_value)

模型注册表集成

注册 PyTorch 模型以进行版本控制和部署

python

import mlflow
import torch.nn as nn
from mlflow import MlflowClient

client = MlflowClient()

with mlflow.start_run():
    # Create a simple model for demonstration
    model = nn.Sequential(
        nn.Conv2d(3, 32, 3),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(32 * 15 * 15, 10),
    )

    # Log model to registry
    model_info = mlflow.pytorch.log_model(
        model, name="pytorch_model", registered_model_name="ImageClassifier"
    )

    # Tag for tracking
    mlflow.set_tags({"model_type": "cnn", "dataset": "imagenet", "framework": "pytorch"})

# Set alias for production deployment
client.set_registered_model_alias(
    name="ImageClassifier",
    alias="champion",
    version=model_info.registered_model_version,
)

分布式训练

跟踪分布式 PyTorch 训练实验

python

import mlflow
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset

# Create synthetic dataset
X_train = torch.randn(1000, 784)
y_train = torch.randint(0, 10, (1000,))
train_dataset = TensorDataset(X_train, y_train)

# Simple model
model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))


def train_epoch(model, train_loader):
    """Simple training epoch for demonstration."""
    model.train()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    total_loss = 0

    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(train_loader)


def train_distributed():
    # Initialize distributed training
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()

    # Wrap model with DDP
    model_ddp = DDP(model.to(rank), device_ids=[rank])

    # Create distributed sampler
    from torch.utils.data.distributed import DistributedSampler

    sampler = DistributedSampler(train_dataset, num_replicas=dist.get_world_size(), rank=rank)
    train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)

    # Only log from rank 0
    if rank == 0:
        mlflow.start_run()
        mlflow.log_params({"world_size": dist.get_world_size(), "backend": "nccl"})

    # Training loop
    epochs = 10
    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # Shuffle data differently each epoch
        train_loss = train_epoch(model_ddp, train_loader)

        # Log metrics from rank 0 only
        if rank == 0:
            mlflow.log_metric("train_loss", train_loss, step=epoch)

    # Save model from rank 0
    if rank == 0:
        mlflow.pytorch.log_model(model, name="distributed_model")
        mlflow.end_run()

了解更多

模型注册表

对 PyTorch 模型进行版本控制和管理

了解更多 →

MLflow Tracking

跟踪实验、参数和指标

了解更多 →