使用 FastAPI 和 Docker 部署机器学习模型分步指南
图片来源：编辑 | Midjourney
您已经训练好了机器学习模型，并且在测试数据上表现出色。但事实是：一个模型放在 Jupyter Notebook 中并不能帮助任何人。只有当您将其部署到生产环境，真正的用户才能从您的工作中受益。
在本文中，我们将使用 scikit-learn 的样本数据集来构建一个糖尿病进展预测器。我们将从原始数据一路构建到一个随时可以部署到云端的容器化 API。
通过跟随本教程进行编码，您将获得：
- 一个已训练的随机森林模型，可预测糖尿病进展分数
- 一个使用 FastAPI 构建的 REST API，接受患者数据并返回预测结果
- 一个已完全容器化、可用于部署的应用程序
让我们开始吧。
设置您的开发环境
在我们开始编码之前，先准备好您的开发环境。您需要：
- Python 3.11+（虽然 3.9+ 也可以）
- 已安装并正在运行的 Docker
- 对 Python 和 API 的基本了解（我会解释非关键部分）
项目结构
这是我们将如何组织项目目录中的所有内容：
|
1 2 3 4 5 6 7 8 9 10 11 12 |
diabetes-predictor/ │ ├── app/ │ ├── __init__.py │ └── main.py # FastAPI 应用程序 │ ├── models/ │ └── diabetes_model.pkl # 训练好的模型（我们将创建它） │ ├── train_model.py # 模型训练脚本 ├── requirements.txt # Python 依赖项 └── Dockerfile # 容器配置 |
安装依赖项
让我们创建一个干净的虚拟环境：
|
1 2 |
$ python -m venv diabetes-env $ source diabetes-env/bin/activate # Windows: diabetes-env\Scripts\activate |
现在安装所需的库：
|
1 |
$ pip install scikit-learn pandas fastapi uvicorn |
构建用于预测糖尿病进展的机器学习模型
让我们从创建机器学习模型开始。创建 train_model.py 文件：
|
1 2 3 4 5 6 7 |
# train_model.py from sklearn.datasets import load_diabetes from sklearn.model_selection import train_test_split from sklearn.ensemble import RandomForestRegressor from sklearn.metrics import mean_squared_error, r2_score import pickle import os |
我们选择随机森林是因为它很稳健，能很好地处理不同的特征尺度，并且还能提供特征重要性洞察。
让我们加载并探索我们的糖尿病数据集：
|
1 2 3 4 5 6 7 |
# 加载糖尿病数据集 diabetes = load_diabetes() X, y = diabetes.data, diabetes.target print(f"数据集形状: {X.shape}") print(f"特征: {diabetes.feature_names}") print(f"目标范围: {y.min():.1f} 到 {y.max():.1f}") |
糖尿病数据集包含 442 名患者的记录，具有 10 个生理特征。目标是基线一年后疾病进展的量化度量：数值越高表示进展越严重。
输出
|
1 2 3 |
数据集 形状: (442, 10) 特征: ['age', 'sex', 'bmi', 'bp', 's1', 's2', 's3', 's4', 's5', 's6'] 目标 范围: 25.0 to 346.0 |
现在我们来准备数据：
|
1 2 3 4 5 6 7 |
# 划分数据 X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42 ) print(f"训练样本数: {X_train.shape[0]}") print(f"测试样本数: {X_test.shape[0]}") |
80/20 的划分给了我们足够的训练数据，同时也保留了可靠的测试集。使用 `random_state=42` 可以确保结果的可复现性。
输出
|
1 2 |
训练 样本数: 353 测试 样本数: 89 |
是时候训练我们的模型了：
|
1 2 3 4 5 6 7 8 |
# 训练随机森林模型 model = RandomForestRegressor( n_estimators=100, random_state=42, max_depth=10 ) model.fit(X_train, y_train) |
我们将 `max_depth` 设置为 10，以防止在该相对较小的数据集上过拟合。拥有 100 棵树，我们在不过度计算的情况下获得了良好的性能。
让我们评估一下我们的模型：
|
1 2 3 4 5 6 7 8 |
# 进行预测并评估 y_pred = model.predict(X_test) mse = mean_squared_error(y_test, y_pred) r2 = r2_score(y_test, y_pred) print(f"均方误差: {mse:.2f}") print(f"R² 分数: {r2:.3f}") |
R² 分数告诉我们模型解释了疾病进展中多大比例的方差。对于这个数据集来说，任何高于 0.4 的分数都算相当不错！
输出
|
1 2 |
均方 误差: 2974.20 R² 分数: 0.439 |
最后，让我们保存训练好的模型：
|
1 2 3 4 5 6 7 |
# 创建 models 目录并保存模型 os.makedirs('models', exist_ok=True) with open('models/diabetes_model.pkl', 'wb') as f: pickle.dump(model, f) print("模型已成功训练并保存！") |
运行此脚本来训练您的模型：
|
1 |
$ python3 train_model.py |
您应该会看到显示模型性能的输出，并确认已保存模型。
创建 FastAPI 应用程序
现在是激动人心的部分：将我们的模型变成一个 Web API。
如果尚未创建，请创建 `app` 目录和空的 `__init__.py` 文件：
|
1 2 |
$ mkdir app $ touch app/__init__.py |
现在创建 `app/main.py` 文件，其中包含我们的 API 代码：
|
1 2 3 4 5 6 |
# app/main.py from fastapi import FastAPI from pydantic import BaseModel import pickle import numpy as np import os |
FastAPI 使用 Pydantic 进行请求验证。这意味着它会自动验证传入数据，并在出现问题时提供清晰的错误消息。
让我们定义输入数据结构：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
# 定义输入数据模式 class PatientData(BaseModel): age: float sex: float bmi: float bp: float # 血压 s1: float # 血清测量 1 s2: float # 血清测量 2 s3: float # 血清测量 3 s4: float # 血清测量 4 s5: float # 血清测量 5 s6: float # 血清测量 6 class Config: schema_extra = { "example": { "age": 0.05, "sex": 0.05, "bmi": 0.06, "bp": 0.02, "s1": -0.04, "s2": -0.04, "s3": -0.02, "s4": -0.01, "s5": 0.01, "s6": 0.02 } } |
示例值有助于 API 用户理解预期的输入格式。请注意，糖尿病数据集的特征已经过标准化。
接下来，我们初始化 FastAPI 应用并将模型加载到 FastAPI 环境中：
|
1 2 3 4 5 6 7 8 9 10 11 |
# 初始化 FastAPI 应用 app = FastAPI( title="糖尿病进展预测器", description="根据生理特征预测糖尿病进展分数", version="1.0.0" ) # 加载训练好的模型 model_path = os.path.join("models", "diabetes_model.pkl") with open(model_path, 'rb') as f: model = pickle.load(f) |
最后，让我们创建预测端点：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
@app.post("/predict") def predict_progression(patient: PatientData): """ 预测糖尿病进展分数 """ # 将输入转换为 numpy 数组 features = np.array([[ patient.age, patient.sex, patient.bmi, patient.bp, patient.s1, patient.s2, patient.s3, patient.s4, patient.s5, patient.s6 ]]) # 进行预测 prediction = model.predict(features)[0] # 返回带附加上下文的结果 return { "predicted_progression_score": round(prediction, 2), "interpretation": get_interpretation(prediction) } def get_interpretation(score): """提供分数的易读解释""" if score < 100: return "低于平均水平的进展" elif score < 150: return "平均水平的进展" else: return "高于平均水平的进展" |
解释函数通过为数值预测提供上下文，有助于使我们的 API 更用户友好。
我们还可以添加一个健康检查端点：
|
1 2 3 |
@app.get("/") def health_check(): return {"status": "healthy", "model": "diabetes_progression_v1"} |
本地测试 API
在容器化之前，让我们在本地测试我们的 API。从项目根目录运行以下命令：
|
1 |
$ uvicorn app.main:app --reload --port 8000 |
在浏览器中打开 `https://:8000/`，您会看到 FastAPI 应用正在运行。尝试使用示例数据进行预测。
您也可以使用 curl 进行测试：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
curl -X POST "https://:8000/predict" \ -H "Content-Type: application/json" \ -d '{ "age": 0.05, "sex": 0.05, "bmi": 0.06, "bp": 0.02, "s1": -0.04, "s2": -0.04, "s3": -0.02, "s4": -0.01, "s5": 0.01, "s6": 0.02 }' |
您应该会得到以下结果：
|
1 2 |
{"predicted_progression_score":213.34, "interpretation":"高于平均水平的进展"} |
使用 Docker 进行容器化
现在我们将所有内容打包成一个 Docker 容器。首先，创建 requirements.txt 文件：
|
1 2 3 4 5 |
fastapi==0.115.12 uvicorn==0.34.2 scikit-learn==1.6.1 pandas==2.2.3 numpy==2.2.6 |
我们固定了特定的版本，以确保跨环境的一致性。
现在创建 Dockerfile：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
# 使用 Python 3.11 slim 镜像 FROM python:3.11-slim # 设置工作目录 WORKDIR /app # 安装系统依赖项（如果需要） RUN apt-get update && apt-get install -y \ && rm -rf /var/lib/apt/lists/* # 复制 requirements 文件并安装 Python 依赖项 COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt # 复制应用程序代码 COPY app/ ./app/ COPY models/ ./models/ # 暴露端口 EXPOSE 8000 # 运行应用程序 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"] |
slim 镜像使我们的容器保持小巧，而 `--no-cache-dir` 可防止 pip 存储缓存的包，进一步减小了大小。
构建您的 Docker 镜像：
|
1 |
$ docker build -t diabetes-predictor . |
运行容器：
|
1 |
$ docker run -d -p 8000:8000 diabetes-predictor |
您的 API 现在正在容器中运行！像之前一样进行测试。
发布到 Docker Hub
现在您的容器化 API 在本地运行正常，让我们通过 Docker Hub 与世界分享它。此步骤对于云部署至关重要。大多数云平台可以直接从 Docker Hub 拉取，使部署无缝进行。
设置 Docker Hub
首先，如果您还没有 Docker Hub 账户，则需要一个。
- 访问 hub.docker.com 并注册。
- 选择一个您喜欢的用户名。它将成为您镜像 URL 的一部分。
登录 Docker Hub
从您的终端登录 Docker Hub：
|
1 |
$ docker login |
您将被提示输入 Docker Hub 用户名和密码。请仔细输入。这将创建一个身份验证令牌，允许您推送镜像。
标记您的镜像
在推送之前，我们需要用您的 Docker Hub 用户名标记我们的镜像。Docker 使用特定的命名约定：
|
1 |
$ docker tag diabetes-predictor your-username/diabetes-predictor:v1.0 |
将 `your-username` 替换为您的实际 Docker Hub 用户名。`v1.0` 是一个版本标签。版本化您的镜像以便跟踪更改和在需要时回滚是一种良好的做法。
我们还创建一个 `latest` 标签，许多部署平台默认使用它：
|
1 |
$ docker tag diabetes-predictor your-username/diabetes-predictor:latest |
检查您的已标记镜像：
|
1 |
$ docker images | grep diabetes-predictor |
您应该会看到三个条目：您的原始镜像和两个新标记的版本。
推送到 Docker Hub
现在让我们将您的镜像推送到 Docker Hub：
|
1 2 |
$ docker push your-username/diabetes-predictor:v1.0 $ docker push your-username/diabetes-predictor:latest |
第一次推送可能需要几分钟，因为 Docker 会上传所有层。之后的推送应该会快很多。
您可以通过拉取并运行已发布的镜像来验证一切是否正常工作。
|
1 2 3 4 5 |
# 首先停止本地容器 $ docker stop $(docker ps -q --filter ancestor=diabetes-predictor) # 从 Docker Hub 拉取并运行 $ docker run -d -p 8000:8000 your-username/diabetes-predictor:latest |
再次测试 API，确保一切仍然正常工作。如果正常，您的模型现在已公开可用，并已准备好进行云部署。
总结
恭喜！您刚刚构建了一个完整的机器学习部署管道。
- 使用医学数据训练了一个强大的随机森林模型
- 使用 FastAPI 创建了一个可工作的 REST API
- 使用 Docker 对应用程序进行了容器化
您的模型现在已准备好进行云部署！您可以将其部署到 AWS ECS、Fargate、Google Cloud 或 Azure。
想进一步扩展？您可以考虑添加以下内容：
- 身份验证和速率限制
- 模型监控和日志记录
- 批量预测端点
您现在已掌握将任何机器学习模型部署到生产环境的基础知识。祝您编码愉快！
令人印象深刻 👏
谢谢您分享这个