如何部署AI模型到生产环境？Flask与Docker实战

2025-06-23

在当今的科技浪潮中，人工智能模型的开发与部署已经成为企业数字化转型的重要组成部分。然而，从实验室到生产环境的过渡并非易事。本文将通过Flask和Docker的实战案例，详细介绍如何高效地将AI模型部署到生产环境中。

一、准备工作

1. 环境搭建

在开始之前，确保你的开发环境已经安装了以下工具：

Python（建议3.8及以上版本）
Flask框架
Docker
AI模型及其依赖库（如TensorFlow、PyTorch等）

2. 模型选择与训练

假设我们已经完成了一个基于TensorFlow的图像分类模型的训练，并保存为model.h5文件。接下来需要将其封装成一个可调用的服务。

二、使用Flask构建API服务

Flask是一个轻量级的Python Web框架，非常适合用于构建小型API服务。以下是具体的步骤：

1. 创建项目结构

创建一个名为ai_service的目录，并在其中添加以下文件：

ai_service/
│
├── app.py          # Flask应用入口
├── model.h5        # 训练好的模型
└── requirements.txt # 依赖项列表

2. 编写Flask应用

在app.py中编写代码以加载模型并提供预测接口：

from flask import Flask, request, jsonify
import tensorflow as tf
import numpy as np
from PIL import Image
import io

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

# 加载模型
model = tf.keras.models.load_model('model.h5')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    try:
        # 读取上传的图片
        img = Image.open(io.BytesIO(file.read()))
        img = img.resize((128, 128))  # 调整图片大小以匹配模型输入
        img_array = np.array(img) / 255.0  # 归一化处理
        img_array = np.expand_dims(img_array, axis=0)

        # 使用模型进行预测
        predictions = model.predict(img_array)
        predicted_class = np.argmax(predictions, axis=1)[0]

        return jsonify({'class': int(predicted_class)})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

3. 定义依赖项

在requirements.txt中列出所有依赖项：

Flask==2.2.2
tensorflow==2.10.0
numpy==1.23.5
Pillow==9.4.0

三、使用Docker容器化应用

为了确保应用能够在不同的环境中一致运行，我们可以使用Docker容器化Flask应用。

1. 编写Dockerfile

在ai_service目录下创建一个名为Dockerfile的文件，内容如下：

# 基础镜像
FROM python:3.8-slim

# 设置工作目录
WORKDIR /app

# 复制当前目录下的文件到容器中的/app目录
COPY . .

# 安装依赖项
RUN pip install --no-cache-dir -r requirements.txt

# 暴露端口
EXPOSE 5000

# 启动Flask应用
CMD ["python", "app.py"]

2. 构建Docker镜像

在终端中进入ai_service目录，并运行以下命令构建镜像：

docker build -t ai-service .

3. 运行Docker容器

启动容器时可以指定端口映射：

docker run -d -p 5000:5000 --name ai-container ai-service

此时，你的AI服务已经在本地的http://localhost:5000/predict地址上运行。

四、测试API接口

可以使用curl命令或Postman来测试API。例如，上传一张图片进行预测：

curl -X POST http://localhost:5000/predict -F "file=@test.jpg"

如果一切正常，你将收到类似以下的JSON响应：

{
    "class": 3
}

五、部署到生产环境

1. 云平台部署

将Docker镜像推送到Docker Hub或其他镜像仓库后，可以轻松地将其部署到AWS、Google Cloud或Azure等云平台上。例如，在AWS ECS中创建任务定义并启动服务。

2. 自动化部署

结合CI/CD工具（如Jenkins、GitHub Actions），可以实现代码更改后的自动构建和部署流程。

六、总结

通过Flask和Docker的组合，我们可以快速将AI模型部署到生产环境中。Flask提供了灵活的API接口设计能力，而Docker则保证了应用的一致性和可移植性。这种方法不仅适用于小型项目，也可以扩展到更复杂的微服务架构中。希望本文能为你在AI模型部署方面提供一些启发和帮助！