使用 MLflow 自动记录 PyTorch Lightning MNIST 训练：从 `mlflow run` 到 `autolog` 全流程实战

使用 MLflow 自动记录 PyTorch Lightning MNIST 训练从mlflow run到autolog全流程实战【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow本文以仓库 examples/pytorch/MNIST 示例为骨架讲解如何用 PyTorch Lightning 训练手写数字识别模型并仅靠一行mlflow.pytorch.autolog()调用自动完成参数、指标与模型含 early stopping 最佳模型的记录。读完本文你将掌握 MLflow Projects 的标准运行方式、训练参数自定义、LightningCLI 直跑脚本以及 autolog 在源码层如何工作能够在自己的 PyTorch Lightning 项目中直接复用这套实验追踪方案。示例概览一行 autolog 替换全部手动埋点该示例的核心思路是代码几乎全部专注于模型训练本身实验追踪只通过一行mlflow.pytorch.autolog()完成。在 mnist_autolog_example.py 中唯一与 MLflow 直接相关的代码是if cli.trainer.global_rank 0: mlflow.pytorch.autolog() cli.trainer.fit(cli.model, datamodulecli.datamodule) cli.trainer.test(ckpt_pathbest, datamodulecli.datamodule)这行调用会自动完成三件事记录参数params如训练轮数epochs、优化器名称与默认参数learning rate 等记录指标metrics每个 epoch 的 loss、验证指标、测试准确率等记录模型models训练完成的模型会被保存为 MLflow 模型 artifact并且由于示例配置了EarlyStopping回调自动记录的是早停后性能最佳的模型而不是最后一次迭代的模型。cli.trainer.test(ckpt_pathbest)则确保用早停保存的最佳 checkpoint 在测试集上评估并将avg_test_acc一并写入该 run。环境准备与依赖示例的依赖声明在两个文件中二选一即可python_env.yaml默认使用MLflow Projects 的推荐方式conda.yaml基于 conda-forge 的旧式环境声明两者声明的核心依赖一致依赖版本要求作用mlflow最新实验追踪与模型管理torch2.1深度学习框架torchvision0.15.1MNIST 数据集与图像变换lightning最新PyTorch Lightning 训练框架jsonargparse[signatures]4.17.0LightningCLI的 CLI 解析依赖protobuf4.0.0与 MLflow 通信兼容在 conda 环境下运行mlflow run .时MLflow 会根据 MLproject 中的python_env: python_env.yaml自动创建隔离环境无需手动pip install。用 MLflow Projects 标准方式运行进入示例目录后直接运行mlflow run .MLflow 会读取 MLproject其内容定义了项目名、环境与入口点参数name: mnist-autolog-example python_env: python_env.yaml entry_points: main: parameters: max_epochs: {type: int, default: 5} devices : {type: str, default: auto} strategy: {type: str, default: auto} accelerator: {type: str, default: auto} batch_size: {type: int, default: 64} num_workers: {type: int, default: 3} learning_rate: {type: float, default: 0.001} command: | python mnist_autolog_example.py \ --trainer.max_epochs{max_epochs} \ --trainer.devices{devices} \ --trainer.strategy{strategy} \ --trainer.accelerator{accelerator} \ --data.batch_size{batch_size} \ --data.num_workers{num_workers} \ --model.learning_rate{learning_rate}注意max_epochs的默认值是5而非脚本层 LightningCLI 的默认值learning_rate默认0.001batch_size默认64。MLflow 会把每个{参数名}替换为实际传入值后执行command参数以--trainer.xxx、--data.xxx、--model.xxx的前缀形式透传给 LightningCLI分别对应 Trainer、DataModule 与 LightningModule 的配置。自定义参数运行需要修改训练轮数时mlflow run . -P max_epochsX其中X是你期望的max_epochs值。-P后可跟任意 MLproject 中声明的参数例如mlflow run . -P max_epochs5 -P devices1 -P batch_size32 -P num_workers2 -P learning_rate0.01 -P strategyddp跳过环境创建直接运行如果你本机已装好所需模块不想让 MLflow 新建隔离环境加--env-managerlocal即可mlflow run . --env-managerlocal该参数适用于任何希望复用当前 Python 环境的场景如 CI 中已预装依赖。训练参数全解README 列出的 7 个参数均可通过命令行覆盖含义如下参数类型/默认值说明max_epochsint默认 5训练轮数训练过程中可按CtrlC提前中断devicesstr默认auto使用的 GPU 数量如1、2auto表示自动探测strategystr默认auto分布式训练策略如ddpDistributed Data Parallel不设置则单进程训练acceleratorstr默认auto硬件加速器如gpuGPU 环境或cpubatch_sizeint默认 64训练输入批大小num_workersint默认 3加载训练数据的工作线程数learning_ratefloat默认 0.001学习率示例内部默认 0.01MLproject 层覆盖为 0.001完整覆盖示例mlflow run . -P max_epochs5 -P devices1 -P batch_size32 -P num_workers2 -P learning_rate0.01 -P strategyddp跳过 MLflow Projects 直接运行脚本也可以绕过 MLflow run直接以 LightningCLI 原生的--组名.参数语法运行 mnist_autolog_example.pypython mnist_autolog_example.py \ --trainer.max_epochs 5 \ --trainer.devices 1 \ --trainer.strategy ddp \ --trainer.accelerator gpu \ --data.batch_size 64 \ --data.num_workers 3 \ --model.learning_rate 0.001这种方式同样会触发mlflow.pytorch.autolog()区别在于不会由 MLflow Projects 管理环境需要你自行确保依赖已安装。注意strategyddp要求devices与可用的 GPU 数量匹配否则 Lightning 会报错。源码走读模型、数据与回调模型三层 MLP 分类器LightningMNISTClassifiermnist_autolog_example.py是一个 784 → 128 → 256 → 10 的三层全连接网络输入展平后的 28×28 像素输出 10 类手写数字的 log-softmax 概率。训练使用F.nll_loss负对数似然优化器为 Adam并配置了ReduceLROnPlateau调度器modemin、factor0.2、patience2、min_lr1e-6监控val_loss动态衰减学习率。数据模块MNIST 三份数据MNISTDataModulemnist_autolog_example.py在setup中自动下载 MNIST 数据集默认缓存到dataset目录将 60000 张训练图按55000 / 5000切分为训练集与验证集测试集保持官方 10000 张。图像预处理使用均值 0.1307、标准差 0.3081 的标准化MNIST 官方统计值。回调早停与最佳模型cli_main中组装了三个 Lightning 回调mnist_autolog_example.pyearly_stopping EarlyStopping(monitorval_loss) checkpoint_callback ModelCheckpoint( dirpathos.getcwd(), save_top_k1, verboseTrue, monitorval_loss, modemin ) lr_logger LearningRateMonitor() cli LightningCLI( LightningMNISTClassifier, MNISTDataModule, runFalse, save_config_callbackNone, trainer_defaults{callbacks: [early_stopping, checkpoint_callback, lr_logger]}, )EarlyStopping监控val_loss验证损失不再下降时提前终止训练ModelCheckpoint(save_top_k1, modemin)只保留val_loss最小的 checkpoint即最佳模型LearningRateMonitor记录每个 epoch 的学习率变化。关键点Lightning 的EarlyStopping只在val_loss有改善时才更新最佳模型因此mlflow.pytorch.autolog()配合这两者记录到的模型正是泛化表现最好的版本——这也是该示例想演示的核心价值。autolog 源码级原理它到底记录了什么mlflow.pytorch.autolog()的实现位于 mlflow/pytorch/init.py通过safe_patch对pl.Trainer.fit打补丁见init.py在fit调用时自动创建 run并向 Trainer 注入一个名为__MlflowPLCallback的内部回调mlflow/pytorch/_lightning_autolog.py。该回调在训练生命周期各节点完成记录on_train_start_lightning_autolog.py记录 tag{Mode: training}、参数{epochs: trainer.max_epochs}、优化器名称及优化器默认参数如lron_train_epoch_end/on_validation_epoch_end_lightning_autolog.py把trainer.callback_metrics中的指标按 epoch 写入 MLflowon_train_batch_end_lightning_autolog.py当设置了log_every_n_step时按 step 记录批级指标并自动跳过 Lightning 训练前的 sanity check 轮次_lightning_autolog.py模型保存训练结束后将最佳 checkpoint 转存为 MLflow 模型若log_model_signaturesTrue会通过 patchmodel.forward捕获首个 batch 的输入输出张量来推断模型签名_lightning_autolog.py。autolog()还提供丰富的配置项可按需在调用时传入参数默认值作用log_every_n_epoch1每 n 个 epoch 记录一次指标log_every_n_stepNone每 n 个 step 记录批级指标设为 1 可能带来性能开销log_modelsTrue是否将训练好的模型保存为 MLflow 模型 artifactlog_datasetsTrue是否记录数据集信息disableFalse是否禁用该集成exclusiveFalseTrue 时 autolog 内容不写入用户手动创建的 runregistered_model_nameNone每次训练后自动注册为新版本模型extra_tagsNone为每个 run 附加自定义 tagcheckpointTrue是否启用自动模型 checkpoint 保存需 pytorch-lightning 1.6.0checkpoint_monitorval_losscheckpoint 监控的指标名checkpoint_modeminmin/max决定最优 checkpoint 的判定方向checkpoint_save_best_onlyTrue只保留最优 checkpointcheckpoint_save_weights_onlyFalseTrue 时仅保存权重否则连同优化器、调度器状态checkpoint_save_freqepochepoch 或整数控制保存频率log_model_signaturesTrue是否记录模型签名注意完整的 autolog 仅对继承LightningModule的模型生效对纯torch.nn.Moduleautolog 只捕获torch.utils.tensorboard.SummaryWriter的add_scalar/add_hparams调用且无 epoch 概念。在 MLflow UI 中查看结果代码执行完毕后启动 MLflow 服务端mlflow server然后在浏览器访问 http://localhost:5000即可看到本次 run 的Parametersepochs、optimizer_name、lr 等Metrics每个 epoch 的val_loss、训练 loss以及测试阶段的avg_test_acc随 epoch 变化的曲线图Artifactsmodel目录下的 MLflow 模型文件、checkpoint 与模型签名。将结果记录到自定义 Tracking Server默认情况下 run 记录在本地mlruns目录。如需写入自定义非默认tracking 服务设置环境变量即可export MLFLOW_TRACKING_URIhttp://localhost:5000/ mlflow run .MLFLOW_TRACKING_URI支持本地路径如file:///...、HTTP 服务地址、Databricks workspace 等多种形式设置后mlflow run与脚本内的mlflow.pytorch.autolog()都会把实验数据写入该地址便于团队共享与多端复现。小结本示例用最少的代码展示了 MLflow 与 PyTorch Lightning 的深度集成MLproject声明参数与环境、mlflow run统一入口、LightningCLI 透传配置、autolog()全自动追踪并借助早停保证记录到的是最佳模型。你可以直接复制 examples/pytorch/MNIST 的目录结构把模型类、数据模块替换为自己的实现即可零成本获得一套完整的实验追踪体系。【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考