
使用 Docker 和 FastAPI 的首次容器化机器学习部署
图片由 Editor | ChatGPT 提供
引言
部署机器学习模型可能看起来很复杂,但现代工具可以简化这一过程。FastAPI 是一个用于构建 API 的高性能 Web 框架,而 Docker 允许你在隔离的容器化环境中运行应用程序。结合这两种技术可以简化跨不同系统的部署,确保可扩展性,并使维护更容易。这种方法有助于避免生产环境中的依赖冲突,从而为服务 ML 模型创建了一个可靠的管道。
在本文中,您将学习如何使用 FastAPI 和 Docker 部署机器学习模型。
准备
在开始之前,请确保你的系统已安装以下软件
- Python 3.8+ – 训练模型和运行 FastAPI 服务器所需的版本
- pip – Python 的包安装程序,用于管理依赖项
- Docker – 一个容器平台,用于在环境中一致地构建和运行应用程序
你还应该熟悉基本的 Python 编程,理解机器学习概念,并熟悉 RESTful API。
这是我们推荐的项目结构
1 2 3 4 5 6 7 8 |
iris-fastapi-app/ ├── app/ │ ├── __init__.py │ └── iris_model.pkl # 训练好的模型 ├── main.py # FastAPI 应用 ├── train_model.py # 训练和保存模型的脚本 ├── requirements.txt # 依赖项 ├── Dockerfile # Docker 构建文件 |
训练机器学习模型
我们将从使用 Scikit-learn 的 Iris 数据集训练一个简单的随机森林分类器开始。下面的脚本(你应该将其保存为 train_model.py
)负责加载数据、训练分类器以及使用 joblib
将模型序列化到文件。根据我们的项目结构,这个保存的模型将被放置在 app/
目录中。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
from sklearn.datasets import load_iris from sklearn.ensemble import RandomForestClassifier import joblib import os def train_and_save_model(): # 确保 'app' 目录存在 os.makedirs('app', exist_ok=True) iris = load_iris() X, y = iris.data, iris.target model = RandomForestClassifier() model.fit(X, y) joblib.dump(model, 'app/iris_model.pkl') print("模型已训练并保存到 app/iris_model.pkl") if __name__ == "__main__": train_and_save_model() |
要训练和保存你的模型,请从终端运行此脚本
1 |
python train_model.py |
创建 FastAPI 应用程序
下一步是通过 API 公开模型,以便其他应用程序或用户可以访问它。FastAPI 使得这个过程变得容易,只需最少的样板代码,并提供出色的类型检查、验证和文档支持。
我们将构建一个简单的 FastAPI 应用程序,该应用程序加载训练好的模型,并提供一个名为 /predict
的端点,以根据用户输入返回预测。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
from fastapi import FastAPI from pydantic import BaseModel import joblib import numpy as np app = FastAPI() model = joblib.load("app/iris_model.pkl") class IrisInput(BaseModel): sepal_length: float sepal_width: float petal_length: float petal_width: float @app.post("/predict") def predict(data: IrisInput): input_data = np.array([[data.sepal_length, data.sepal_width, data.petal_length, data.petal_width]]) prediction = model.predict(input_data) return {"prediction": int(prediction[0])} |
该应用公开了一个名为 /predict
的端点,它接受花朵的测量值并返回预测的类别。
编写 Dockerfile
要在容器化环境中运行此 FastAPI 应用程序,你需要创建一个 Dockerfile。此文件包含 Docker 构建镜像的说明,该镜像打包了你的应用程序及其依赖项。在项目根目录创建一个文件,内容如下,并将其命名为 Dockerfile
,不带文件扩展名。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
# 使用官方 Python 运行时作为基础镜像 FROM python:3.10-slim # 设置工作目录 WORKDIR /app # 复制 requirements 并安装依赖项 COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt # 复制应用程序的其余代码 COPY . . # 暴露端口 EXPOSE 8000 # 运行应用程序 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] |
创建 requirements.txt 文件
requirements.txt
文件用于 pip 在你的 Docker 容器中安装所有必需的依赖项。它应包含项目中使用的所有库
1 2 3 4 5 |
fastapi uvicorn scikit-learn joblib numpy |
你可以手动创建此文件,或者通过运行以下命令生成
1 |
pip freeze > requirements.txt |
构建和运行 Docker 容器
一旦你的 FastAPI 应用程序、模型和 Dockerfile 都准备就绪,下一步就是使用 Docker 将应用程序容器化,然后运行它。此过程可确保你的应用程序能在任何环境中可靠地运行。
首先,构建 Docker 镜像
1 |
docker build -t iris-fastapi-app . |
然后,运行容器
1 |
docker run -d -p 8000:8000 iris-fastapi-app |
此命令将容器的 8000 端口映射到你的本地 8000 端口。
测试 API 端点
现在你的 FastAPI 应用已在 Docker 容器中运行,你可以本地测试该 API。
你可以使用浏览器或 Postman 等工具进行验证,或者也可以使用 curl
1 2 3 |
curl -X POST "https://:8000/predict" \ -H "Content-Type: application/json" \ -d '{"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2}' |
预期输出
1 |
{"prediction": 0} |
FastAPI 还提供了交互式文档,地址为 https://:8000/docs
。你可以使用这个 Swagger UI 直接在浏览器中测试和调试 /predict
端点。
改进模型服务
虽然基本设置在初始部署中效果很好,但实际场景通常需要改进来提升开发体验和管理特定环境的配置。以下是一些需要注意的要点。
在开发期间启用实时重新加载
在本地开发时,启用自动重新加载会很有帮助,这样你的 API 在你更改代码时就能重新启动。FastAPI 使用 Uvicorn,它支持此功能。在你的 Dockerfile
中修改 CMD
行,或者在开发时通过 docker-compose
使用它,如下所示
1 |
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"] |
注意:--reload
标志仅用于开发。在生产环境中避免使用它。
使用环境变量
不要硬编码路径或配置值,而是使用环境变量。这使得你的应用程序更灵活、更具生产力。
例如,你可以在 main.py
中重构模型加载
1 2 3 4 |
import os model_path = os.getenv("MODEL_PATH", "app/iris_model.pkl") model = joblib.load(model_path) |
你也可以在 Docker run 命令中传递环境变量
1 |
docker run -d -p 8000:8000 -e MODEL_PATH=app/iris_model.pkl iris-fastapi-app |
结论
使用 FastAPI 和 Docker 部署机器学习模型是一种高效且可扩展的方法。FastAPI 提供了一种高性能的方法将模型公开为 API,而 Docker 则确保了跨所有环境的一致行为。两者结合起来,创建了一个强大的工作流程,简化了开发、测试和部署,并帮助你的机器学习服务变得更加健壮并为生产做好准备。
暂无评论。