# 部署深度学习模型：Flask API 服务端与客户端通信实战

www_pp_

已于 2025-05-06 15:46:30 修改

阅读量788

点赞数 13

文章标签：深度学习 flask 人工智能

于 2025-05-05 19:45:26 首次发布

本文链接：https://blog.csdn.net/www_pp_/article/details/147721519

版权

部署深度学习模型：Flask API 服务端与客户端通信实战

在这篇文章中，我们将探讨如何使用 Flask 框架部署一个深度学习模型，并通过客户端与服务端进行通信。我们将通过一个实际的例子，展示如何构建服务端和客户端，以及如何处理图像预测请求。
在这里插入图片描述

环境准备

首先，确保你已经安装了以下库：

Flask
PyTorch
torchvision
Pillow

如果尚未安装，可以通过以下命令安装：

pip install flask torch torchvision pillow

服务端代码

服务端代码的主要功能是加载预训练的深度学习模型，接收客户端发送的图像数据，进行预测，并将结果返回给客户端。

# 导入所需的库
import io  # 用于处理二进制数据
import flask  # Flask框架，用于搭建Web服务
import torch  # PyTorch库，用于深度学习模型的加载和推理
import torch.nn.functional as F  # PyTorch的神经网络函数模块，用于softmax等操作
from PIL import Image  # Python图像处理库，用于图像的读取和预处理
from torch import nn  # PyTorch的神经网络模块
from torchvision import transforms, models  # torchvision库，用于图像预处理和加载预训练模型

# 初始化Flask app
app = flask.Flask(__name__)  # 创建一个新的Flask应用程序实例
# __name__参数通常被传递给Flask应用程序来定位应用程序的根路径，这样Flask就可以知道在哪里找到模板、静态文件等。
# 总体来说app = flask.Flask(__name__)是FLasK应用程序的起点。它初始化了一个新的FLasK应用程序实例。
model = None  # 初始化模型变量为None
use_gpu = False  # 初始化是否使用GPU的标志为False

# 定义加载模型的函数
def load_model():
    # """Load the pre-trained model, you can use your model just as easily."""
    global model  # 声明使用全局变量model
    # 加载resnet18网络
    model = models.resnet18()  # 加载预训练的resnet18模型
    num_ftrs = model.fc.in_features  # 获取全连接层的输入特征数量
    model.fc = nn.Sequential(nn.Linear(num_ftrs, 102))  # 修改全连接层，输出为102个类别（根据具体任务修改类别数）

    # print(model)
    checkpoint = torch.load('best.pth')  # 加载训练好的模型权重
    model.load_state_dict(checkpoint['state_dict'])  # 将权重加载到模型中
    # 将模型指定为测试格式
    model.eval()  # 将模型设置为评估模式

    # 是否使用gpu
    if use_gpu:
        model.cuda()  # 如果使用GPU，则将模型移动到GPU上

# 定义数据预处理函数
def prepare_image(image, target_size):
    """
    Do image preprocessing before prediction on any data.
    param image    :     original image
    param target_size     :     target image size
    return     :     preprocessed image
    """
    # 针对不同模型，image的格式不同，但需要统一到RGB格式
    if image.mode != 'RGB':
        image = image.convert('RGB')  # 如果图像不是RGB格式，则转换为RGB格式

    # Resize the input image and preprocess it.(按照所使用的模型将输入图片的尺寸修改)
    image = transforms.Resize(target_size)(image)  # 调整图像大小为目标尺寸
    image = transforms.ToTensor()(image)  # 将图像转换为Tensor

    # Convert to Torch. Tensor and normalize. mean与std
    image = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)  # 对图像进行标准化处理

    # Add batch size axis 增加一个维度，用于按batch测试
    image = image[None]  # 增加一个维度
    if use_gpu:
        image = image.cuda()  # 如果使用GPU，则将图像移动到GPU上
    return torch.tensor(image)  # 返回预处理后的图像Tensor

# 定义一个装饰器，用于将指定的URL路径与一个函数关联起来，并指定该函数响应的HTTP方法
@app.route('/predict', methods=['POST'])
def predict():  # 当客户端发送请求时
    # 做一个标志，刚开始无图像传入时为false，传入图像时为true
    data = {'success': False}  # 初始化返回数据字典，初始值为False

    if flask.request.method == 'POST':  # 如果收到POST请求
        if flask.request.files.get('image'):  # 判断是否有图像文件
            image = flask.request.files['image'].read()  # 将收到的图像进行读取，内容为二进制
            image = Image.open(io.BytesIO(image))  # 将二进制图像数据转换为PIL图像对象

            # 利用上面的预处理函数将读入的图像进行预处理
            image = prepare_image(image, target_size=(224, 224))  # 对图像进行预处理，目标尺寸为224x224

            preds = F.softmax(model(image), dim=1)  # 得到各个类别的概率
            results = torch.topk(preds.cpu().data, k=3, dim=1)  # 概率最大的前3个结果
            # torch.topk用于返回输入张量中每行最大的k个元素及其对应的索引
            results = (results[0].cpu().numpy(), results[1].cpu().numpy())  # 将结果转换为numpy数组

            # 将data字典增加一个key，value，其中value为list格式
            data['prediction'] = list()  # 初始化预测结果列表

            for prob, label in zip(results[0][0], results[1][0]):  # 遍历概率和标签
                # Label name = idx2labellstr(label)]
                r = {'label': str(label), 'probability': float(prob)}  # 创建一个字典，包含标签和概率
                # 将预测结果添加至data字典
                data['prediction'].append(r)  # 将预测结果添加到列表中

            # Indicate that the request was a success.
            data['success'] = True  # 将请求成功的标志设置为True
    return flask.jsonify(data)  # 返回预测结果的JSON格式数据

# 主程序入口
if __name__ == '__main__':  # 判断是否是主程序运行
    print('Loading PyTorch model and Flask starting server ...')  # 打印加载模型和启动服务器的信息
    print('Please wait until server has fully started')  # 提示用户等待服务器启动完成

    load_model()  # 先加载模型
    # 再开启服务
    app.run(port='5012')  # 启动Flask应用，监听5012端口

客户端代码

客户端代码负责发送图像数据到服务端，并接收预测结果。

import requests  # 导入requests库，用于发送HTTP请求

# 定义Flask服务的URL地址
flask_url = 'http://127.0.0.1:5012/predict'  # Flask服务的地址，运行在本地主机的5012端口

# 定义一个函数，用于发送图像到Flask服务并获取预测结果
def predict_result(image_path):
    # 打开图像文件并读取其内容
    image = open(image_path, 'rb').read()  # 以二进制模式打开图像文件并读取内容
    payload = {'image': image}  # 将图像内容封装为一个字典，作为请求的文件数据

    # 使用requests库发送POST请求到Flask服务
    r = requests.post(flask_url, files=payload).json()  # 发送POST请求，并将返回的JSON数据解析为字典

    # 检查请求是否成功
    if r['success']:
        # 如果请求成功，遍历预测结果并打印
        for (i, result) in enumerate(r['prediction']):
            print('{}. 预测类别为{}的概率:{}'.format(i + 1, result['label'], result['probability']))
    else:
        # 如果请求失败，打印失败信息
        print('request failed')

# 主程序入口
if __name__ == '__main__':
    # 定义要预测的图像路径
    image_path = r'D:\Users\妄生\PycharmProjects\人工智能\深度学习\模型部署\flower_data\flower_data\val_filelist\image_00059.jpg'
    # 调用predict_result函数，对指定图像进行预测
    predict_result(image_path)