企业官网建设流程全解析

PyTorch Lightning 模型生产部署基础指南：checkpoint 加载、predict_step 与分布式推理

【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

导读

本文是 PyTorch Lightning 生产部署（Deploy）入门系列的基础篇，面向所有使用 Lightning 训练模型的开发者。你将学会三种把训练好的模型投入推理/预测的核心手段：用load_from_checkpoint快速加载权重、用predict_step+Trainer.predict消除预测样板代码、用BasePredictionWriter在多卡（如 DDP）场景下自由落地预测结果。阅读完本文，你可以直接照搬文中代码，把任意 Lightning 模型改造成可批量、可分布式运行的推理管线。

本文对应的原始文档位于仓库 docs/source-pytorch/deploy/production_basic.rst，文中涉及的核心源码分别位于src/lightning/pytorch/core/module.py、src/lightning/pytorch/trainer/trainer.py与src/lightning/pytorch/callbacks/prediction_writer.py。

一、加载 checkpoint 进行预测：load_from_checkpoint

训练完成后，最直接的推理方式是通过LightningModule的类方法load_from_checkpoint加载权重，然后手动进入推理模式执行前向：

model = LitModel.load_from_checkpoint("best_model.ckpt") model.eval() x = torch.randn(1, 64) with torch.no_grad(): y_hat = model(x)

1.1 方法签名与参数说明

从源码 src/lightning/pytorch/core/module.py 可以看到该方法的完整签名：

@classmethod def load_from_checkpoint( cls, checkpoint_path: Union[_PATH, IO], map_location: _MAP_LOCATION_TYPE = None, hparams_file: Optional[_PATH] = None, strict: Optional[bool] = None, weights_only: Optional[bool] = None, **kwargs: Any, ) -> Self:

各参数的作用如下：

1.2 典型用法

# 直接加载（超参数已随 checkpoint 保存） model = MyLightningModule.load_from_checkpoint("path/to/checkpoint.ckpt") # 跨设备映射：GPU 1 → GPU 0 model = MyLightningModule.load_from_checkpoint( "path/to/checkpoint.ckpt", map_location={"cuda:1": "cuda:0"}, ) # 权重与超参数分离存储 model = MyLightningModule.load_from_checkpoint( "path/to/checkpoint.ckpt", hparams_file="path/to/hparams_file.yaml", ) # 覆盖部分超参数后重新初始化 model = MyLightningModule.load_from_checkpoint( PATH, num_layers=128, pretrained_ckpt_path=NEW_PATH, ) # 推理 pretrained_model.eval() pretrained_model.freeze() y_hat = pretrained_model(x)

1.3 使用注意事项（源码明确提示）

从 module.py 的文档字符串中可以提取出三条重要约束：

1. load_from_checkpoint是类方法，必须用LightningModule的类来调用，而不是实例，否则会抛出TypeError；
2. 若模型重写了configure_model钩子，加载时会自动调用它以保证所有层都能从 checkpoint 恢复；
3. 该方法不支持加载分片（sharded）checkpoint，超大模型可能内存不足，此时应改用Trainer.fit(ckpt_path=...)的方式加载。

二、用 LightningModule 的 predict_step 消除预测样板代码

手动加载 checkpoint 再写循环，仍然需要自己管理预测 epoch 的样板代码。LightningModule提供的predict step机制把这份样板代码收进了框架内部。

2.1 最小实现

class MyModel(LightningModule): def predict_step(self, batch, batch_idx, dataloader_idx=0): return self(batch)

然后任意 dataloader 都可以直接交给 Lightning Trainer：

data_loader = DataLoader(...) model = MyModel() trainer = Trainer() predictions = trainer.predict(model, data_loader)

2.2 默认行为与钩子机制

从源码 src/lightning/pytorch/core/module.py 可以看到，基类LightningModule.predict_step的默认实现是：

def predict_step(self, *args: Any, **kwargs: Any) -> Any: # For backwards compatibility batch = kwargs.get("batch", args[0]) return self(batch)

也就是说，不重写predict_step时，预测默认等价于调用self(batch)（即forward）。重写该方法的目的是加入额外的处理逻辑；同时predict_step正是 Lightning 在多设备上扩展推理的核心入口，配合Trainer.predict使用即可获得批量化、分布式的推理能力。

钩子签名中的三个参数含义如下：

• batch：数据迭代器的输出，通常是torch.utils.data.DataLoader的一个批次；
• batch_idx：当前批次的索引；
• dataloader_idx：产生该批次的数据加载器索引（仅在传入多个 dataloader 时使用）。

2.3Trainer.predict的参数

Trainer.predict定义于 src/lightning/pytorch/trainer/trainer.py，其完整签名支持以下常用参数：

返回值为一个列表，其中每个元素对应一个 dataloader 的全部预测结果。

2.4 预测循环的内部结构

Trainer.predict背后由预测循环_PredictionLoop驱动，实现位于 src/lightning/pytorch/loops/prediction_loop.py。从源码结构可以推断出几点：

• 循环内部以dataloaders × batches × samples的层级组织数据，并通过epoch_batch_indices、current_batch_indices记录每个批次对应的原始数据索引，供PredictionWriter使用；
• predictions属性在只有单个 dataloader 时返回第一个 dataloader 的预测结果，多 dataloader 时返回按 dataloader 分组的列表；
• return_predictions的 setter 会检测策略是否使用了_MultiProcessingLauncher：spawn/fork 类策略下若坚持要求返回预测结果，会直接抛出MisconfigurationException——这正是分布式推理时要用BasePredictionWriter落盘的原因。

三、在 predict_step 中实现复杂预测逻辑

当数据需要复杂的预处理或后处理时，把这些逻辑放进predict_step即可。文档以Monte Carlo Dropout为例：预测时保持 dropout 开启并多次前向取平均，从而得到带不确定性的预测。

class LitMCdropoutModel(L.LightningModule): def __init__(self, model, mc_iteration): super().__init__() self.model = model self.dropout = nn.Dropout() self.mc_iteration = mc_iteration def predict_step(self, batch, batch_idx): # enable Monte Carlo Dropout self.dropout.train() # take average of `self.mc_iteration` iterations pred = [self.dropout(self.model(x)).unsqueeze(0) for _ in range(self.mc_iteration)] pred = torch.vstack(pred).mean(dim=0) return pred

要点拆解：

1. 将 dropout 切到训练模式：self.dropout.train()确保每次前向都会随机丢弃神经元，这是蒙特卡洛采样的前提；
2. 多次前向采样：循环self.mc_iteration次，每次把结果unsqueeze(0)堆叠成(mc_iteration, ...)形状；
3. 取均值：torch.vstack(pred).mean(dim=0)得到所有采样结果的均值作为最终预测。

这种模式表明predict_step并不只是forward的替身——你可以在其中自由组合模型、采样、聚合等任意逻辑，且这些逻辑在多卡推理时由 Lightning 自动并行调度。

四、开启分布式推理：BasePredictionWriter

使用predict_step后，Lightning 借助BasePredictionWriter回调为你提供免费的多卡分布式推理能力——每个进程只负责自己 rank 的数据，预测结果按 rank 落盘，无需手动做进程间通信。

4.1 按 epoch 结束写盘（多卡示例）

import torch from lightning.pytorch.callbacks import BasePredictionWriter class CustomWriter(BasePredictionWriter): def __init__(self, output_dir, write_interval): super().__init__(write_interval) self.output_dir = output_dir def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices): # this will create N (num processes) files in `output_dir` each containing # the predictions of it's respective rank torch.save(predictions, os.path.join(self.output_dir, f"predictions_{trainer.global_rank}.pt")) # optionally, you can also save `batch_indices` to get the information about the data index # from your prediction data torch.save(batch_indices, os.path.join(self.output_dir, f"batch_indices_{trainer.global_rank}.pt")) # or you can set `write_interval="batch"` and override `write_on_batch_end` to save # predictions at batch level pred_writer = CustomWriter(output_dir="pred_path", write_interval="epoch") trainer = Trainer(accelerator="gpu", strategy="ddp", devices=8, callbacks=[pred_writer]) model = BoringModel() trainer.predict(model, return_predictions=False)

4.2write_interval的三种取值

BasePredictionWriter的源码 src/lightning/pytorch/callbacks/prediction_writer.py 中定义了WriteInterval枚举，对应三种写盘时机：

若传入非法取值，构造器会抛出MisconfigurationException（错误信息为`write_interval` should be one of ['batch', 'epoch', 'batch_and_epoch']），这一点由单元测试 tests/tests_pytorch/callbacks/test_prediction_writer.py 明确验证。

4.3 按 batch 写盘

当你需要逐 batch 及时落盘（例如推理流式输出、显存压力大时）时，可以设置write_interval="batch"并重写write_on_batch_end：

class CustomWriter(BasePredictionWriter): def __init__(self, output_dir, write_interval): super().__init__(write_interval) self.output_dir = output_dir def write_on_batch_end( self, trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx ): torch.save(prediction, os.path.join(self.output_dir, str(dataloader_idx), f"{batch_idx}.pt"))

4.4 底层调度机制

BasePredictionWriter本质上是Callback的子类，它通过重写两个 Trainer 钩子来驱动用户自定义的写盘方法（见 prediction_writer.py）：

• on_predict_batch_end：若interval.on_batch为真，则从trainer.predict_loop.current_batch_indices取出当前批次的原始数据索引，调用write_on_batch_end；
• on_predict_epoch_end：若interval.on_epoch为真，则把trainer.predict_loop.epoch_batch_indices与trainer.predict_loop.predictions一起传给write_on_epoch_end。

由此可以确认两个关键行为：

1. batch_indices不是凭空而来，而是由预测循环在 prediction_loop.py 中逐批累积维护的，用于把预测结果与原始数据的行号对齐；
2. 在 spawn/fork 类策略（ddp_spawn、多卡 TPU）下，Trainer.predict无法返回预测结果，官方建议必须配合BasePredictionWriter把预测写入磁盘或数据库——这正是文档所说“免费分布式推理”的完整实现方式。

4.5 测试验证

仓库测试 tests/tests_pytorch/callbacks/test_prediction_writer.py 对三种write_interval的触发次数做了严格断言：

• "batch_and_epoch"：write_on_batch_end调用 4 次（与 batch 数一致）、write_on_epoch_end调用 1 次，无论是否设置return_predictions=False；
• "batch"：只触发write_on_batch_end，epoch 回调调用 0 次；
• "epoch"：只触发write_on_epoch_end，batch 回调调用 0 次。

同时 test_prediction_writer.py 还验证了batch_indices的传递：例如batch_size=4、limit_predict_batches=4时，write_on_batch_end收到的批次索引依次是[0,1,2,3]、[4,5,6,7]、[8,9,10,11]、[12,13,14,15]。此外，tests/tests_pytorch/trainer/test_trainer.py（约 L1288 起）也提供了CustomPredictionWriter与Trainer.predict配合的端到端示例，可作为你编写自定义 Writer 的参考模板。

五、生产部署要点小结

把上述内容串起来，一个“从 checkpoint 到分布式推理”的完整生产链路如下：

1. 训练产出 checkpoint：训练过程中由ModelCheckpoint回调保存best_model.ckpt（其中同时包含权重与hyper_parameters）；
2. 加载模型：用MyLightningModule.load_from_checkpoint(...)恢复模型与超参数，必要时用map_location做设备映射、用**kwargs覆盖超参数；
3. 定义predict_step：默认等价于forward；需要复杂前后处理（如 Monte Carlo Dropout）时重写它；
4. 选择推理规模：单机单卡直接用trainer.predict(model, data_loader)拿返回值；多卡/多进程场景设置accelerator、strategy、devices，并挂载BasePredictionWriter回调、令return_predictions=False，让每个 rank 把属于自己的预测与batch_indices写入独立文件。

这三层能力（load_from_checkpoint→predict_step→BasePredictionWriter）构成了 PyTorch Lightning 生产推理的基础栈。如果你需要更深入的内容——例如推理服务化、模型导出（TorchScript/ONNX）、以及大规模 checkpoint 的高效加载——可以继续阅读仓库中同一目录下的进阶文档 docs/source-pytorch/deploy/production_advanced.rst 与 docs/source-pytorch/deploy/production_advanced_2.rst。

【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning