从实验室到生产：用PyTorch Lightning + Flask快速部署你的AI模型（保姆级教程）

当你完成了一个PyTorch Lightning模型的训练，看着验证集上的指标节节攀升，那种成就感无与伦比。但接下来呢？如何让这个精心调教的模型走出实验室，真正为他人所用？本文将带你走过从.ckpt文件到可调用API服务的完整旅程，解决模型部署"最后一公里"的问题。

想象一下：你的图像分类模型可以接收用户上传的照片并实时返回预测结果；或者你的文本情感分析API能被集成到客服系统中。这些场景的实现并不需要复杂的工程知识，只需PyTorch Lightning的基础知识和一些简单的Web开发技巧。我们将使用Flask构建轻量级API，并通过Docker实现环境隔离和便捷部署。

1. 模型准备与权重加载

在开始部署前，确保你的PyTorch Lightning模型已经训练完成并保存为.ckpt文件。这个文件包含了模型架构和训练好的权重，是我们部署的基础。不同于常规PyTorch的.pth文件，.ckpt文件还额外保存了训练时的超参数和优化器状态。

加载Lightning模型需要特别注意其特有的封装方式。以下是一个标准的加载流程：

python复制import torch
from your_model_module import YourLightningModel

# 加载预训练模型
model = YourLightningModel.load_from_checkpoint(
    "path/to/your_model.ckpt",
    input_dim=128,  # 必须与训练时参数一致
    output_dim=10
)
model.eval()  # 切换到推理模式

常见问题排查：

• 如果遇到"Missing class attributes"错误，检查load_from_checkpoint时是否传入了所有必需的初始化参数
• 确保推理环境的PyTorch和PyTorch Lightning版本与训练环境一致
• 对于自定义的LightningModule，需要确保类定义在可导入的模块中

提示：在生产环境中，建议将模型加载包装在异常处理中，并添加加载超时机制，避免服务启动时因模型加载问题而崩溃。

2. 构建Flask API服务

Flask是一个轻量级的Python Web框架，非常适合快速构建模型API。我们将创建一个能够同时处理图像和文本输入的通用API结构。

首先建立基本的应用结构：

code复制/flask_api
│── app.py          # 主应用文件
│── model_loader.py # 模型加载模块
│── utils/          # 辅助函数
│   └── preprocess.py
└── requirements.txt

app.py的核心内容如下：

python复制from flask import Flask, request, jsonify
import torch
from model_loader import load_model

app = Flask(__name__)
model = load_model()  # 初始化时加载模型

@app.route('/predict', methods=['POST'])
def predict():
    try:
        data = request.get_json()
        inputs = preprocess_data(data)  # 数据预处理
        with torch.no_grad():
            outpu