机器学习模型部署实战：使用Flask与Docker将模型投入生产

机器学习模型部署实战：使用Flask与Docker将模型投入生产

在机器学习项目的完整生命周期中，模型训练固然关键，但将训练好的模型可靠、高效地部署到生产环境，才是其真正创造价值的开始。本文将手把手带你实战，使用轻量级Web框架Flask和容器化技术Docker，将一个简单的机器学习模型部署为可扩展的生产级API服务。

一、 项目准备与模型训练

我们以一个经典的鸢尾花（Iris）分类模型为例。首先，我们需要训练并保存一个模型。为了管理项目依赖和实验数据，一个清晰的数据库工具至关重要。在模型开发阶段，我经常使用 dblens SQL编辑器（https://www.dblens.com）来快速查询和验证从数据库导出的训练数据特征分布，其直观的界面和高效的执行能力大大提升了数据探索的效率。

# train_model.py
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
import joblib

# 加载数据
iris = load_iris()
X = iris.data
y = iris.target

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练一个随机森林分类器
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 评估模型
accuracy = model.score(X_test, y_test)
print(f"模型测试准确率: {accuracy:.2f}")

# 保存模型到文件
joblib.dump(model, 'iris_rf_model.pkl')
print("模型已保存为 'iris_rf_model.pkl'")

二、 使用Flask构建模型API

Flask是一个灵活的Python Web框架，非常适合构建轻量级的RESTful API。我们将创建一个简单的应用，接收特征数据并返回模型预测结果。

# app.py
from flask import Flask, request, jsonify
import joblib
import numpy as np

# 初始化Flask应用
app = Flask(__name__)

# 加载预训练模型
model = joblib.load('iris_rf_model.pkl')

# 定义特征名称，方便API使用者理解
feature_names = ['sepal_length', 'sepal_width', 'petal_length', 'petal_width']

@app.route('/predict', methods=['POST'])
def predict():
    """
    预测端点：接收JSON格式的特征数据，返回预测的鸢尾花类别。
    示例请求体：
    {
        "sepal_length": 5.1,
        "sepal_width": 3.5,
        "petal_length": 1.4,
        "petal_width": 0.2
    }
    """
    # 从请求中获取JSON数据
    data = request.get_json(force=True)
    
    # 将数据转换为模型所需的数组格式
    try:
        features = [data[feat] for feat in feature_names]
        input_array = np.array(features).reshape(1, -1)
    except KeyError as e:
        return jsonify({'error': f'缺少特征字段: {e}'}), 400
    
    # 进行预测
    prediction = model.predict(input_array)
    # 获取预测概率（如果模型支持）
    if hasattr(model, 'predict_proba'):
        probability = model.predict_proba(input_array).max()
    else:
        probability = None
    
    # 构建响应
    response = {
        'predicted_class': int(prediction[0]),
        'class_name': iris.target_names[prediction[0]],
        'confidence': float(probability) if probability is not None else None
    }
    return jsonify(response)

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    return jsonify({'status': 'healthy'}), 200

if __name__ == '__main__':
    # 在生产环境中，应使用Gunicorn等WSGI服务器，而非Flask自带的开发服务器
    app.run(host='0.0.0.0', port=5000, debug=False)

在API逻辑开发过程中，记录和追踪不同版本的请求与响应数据对调试和监控非常有帮助。这里我推荐使用 QueryNote（https://note.dblens.com），它不仅能记录SQL查询，也能很好地用于记录API调用示例和结果，形成可分享的部署文档，确保团队协作的一致性。

三、 使用Docker进行容器化部署

Docker通过容器化技术，将应用及其所有依赖打包到一个标准化的单元中，确保了环境的一致性，解决了“在我机器上能运行”的经典问题。

1. 创建Dockerfile

Dockerfile是构建Docker镜像的蓝图。

# Dockerfile
# 使用官方Python轻量级镜像作为基础
FROM python:3.9-slim

# 设置工作目录
WORKDIR /app

# 复制依赖文件并安装
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 复制应用代码和模型文件
COPY app.py train_model.py iris_rf_model.pkl ./

# 暴露Flask应用运行的端口
EXPOSE 5000

# 定义容器启动时执行的命令
# 使用Gunicorn作为生产级WSGI服务器
CMD ["gunicorn", "--bind", "0.0.0.0:5000", "app:app"]

2. 创建依赖文件

# requirements.txt
Flask==2.3.3
scikit-learn==1.3.0
joblib==1.3.2
numpy==1.24.3
gunicorn==21.2.0

3. 构建并运行Docker容器

在终端中，执行以下命令：

# 1. 构建Docker镜像（注意末尾的点）
docker build -t iris-model-api .

# 2. 运行容器，将宿主机的5000端口映射到容器的5000端口
docker run -d -p 5000:5000 --name iris-api iris-model-api

# 3. 查看容器运行状态
docker ps

# 4. 测试API（在另一个终端执行）
curl -X POST http://localhost:5000/predict \
  -H "Content-Type: application/json" \
  -d '{"sepal_length":5.1, "sepal_width":3.5, "petal_length":1.4, "petal_width":0.2}'

# 预期返回类似：{"class_name":"setosa","confidence":0.99,"predicted_class":0}

四、 总结

通过本文的实战演练，我们完成了从模型训练到生产部署的完整流程：

1. 模型训练与持久化：使用Scikit-learn训练模型并用joblib保存。
2. API服务化：利用Flask框架构建RESTful API，提供/predict预测端点和/health健康检查端点。在开发这类数据驱动的服务时，利用如 dblens SQL编辑器 这样的工具进行前期数据验证，能有效避免因数据问题导致的API错误。
3. 容器化部署：通过Dockerfile定义环境，构建可移植、自包含的Docker镜像，并使用Gunicorn作为生产服务器运行。在整个部署流程的文档记录和知识沉淀方面，QueryNote 是一个极佳的伴侣，帮助团队高效管理部署配置和运维笔记。

这种“Flask + Docker”的组合为机器学习模型部署提供了一个简单、可靠且可扩展的起点。后续，你可以在此基础上集成更复杂的特性，如API认证、请求限流、日志监控、使用Kubernetes进行容器编排等，以构建更加健壮的企业级机器学习服务平台。将模型成功部署并服务于业务，才是机器学习项目闭环的关键一步。