Flask + PyTorch模型部署实战：从训练权重到API接口完整工程解析（附完整代码）

梦想三三2026-07-07 10:27

前言

在深度学习项目的落地环节中，模型训练只是万里长征的第一步。如何将训练好的模型权重封装成可调用的Web服务，让其他应用程序（如移动端、前端页面、第三方系统）能够便捷地使用模型推理能力，是算法工程师必须掌握的工程技能。

然而，很多初学者在完成模型训练后，面对模型部署这一环节常常感到无从下手：Flask服务如何搭建？接口路由如何设计？图片数据如何传输和预处理？预测结果如何返回？部署过程中遇到的各种报错又该如何调试？

本文将以一个ResNet18花卉分类模型（102类） 为例，从零开始构建一套完整的PyTorch模型Flask部署解决方案。你将看到如何将训练好的权重文件（best.pth）加载到Flask服务中、如何设计图像预处理流水线、如何编写核心预测接口/predict、如何通过GET/POST测试接口辅助调试，以及如何使用Python脚本、curl命令、Postman等多种方式调用接口进行验证。整篇代码遵循工程化规范，接口统一封装标准JSON返回体，部署流程清晰可复现，可直接迁移至其他图像分类模型的部署任务中。

文章目录

一、完整代码展示

服务端代码（flask_predict.py）

python 复制代码 
import io
import flask
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
from torchvision import transforms, models

# 1. 初始化Flask应用
app = flask.Flask(__name__)
model = None
use_gpu = False
# 类别映射字典（根据你的花卉数据集自行补全）
idx2label = {}


# 2. 加载预训练图像分类模型
def load_model():
    global model, use_gpu
    # 示例：加载resnet预训练模型，替换为你训练好的权重
    model = models.resnet18(weights=None)
    # 此处替换成你的模型权重路径 model.load_state_dict(torch.load("flower_model.pth"))
    model.eval()  # 模型设为推理模式，关闭dropout等训练层
    # 判断是否可用GPU
    if torch.cuda.is_available():
        use_gpu = True
        model.cuda()


# 3. 图片预处理函数
def prepare_image(image, target_size):
    # 构建图像预处理流水线：缩放→转张量→归一化
    transform = transforms.Compose([
        transforms.Resize(target_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    img_tensor = transform(image)
    # 增加batch维度 [C,H,W] → [1,C,H,W]
    img_tensor = torch.unsqueeze(img_tensor, 0)
    if use_gpu:
        img_tensor = img_tensor.cuda()
    return img_tensor


# 4. 核心预测接口 /predict，仅接收POST请求
@app.route("/predict", methods=["POST"])
def predict():
    # 初始化返回数据，默认失败
    data = {"success": False}
    # 判断请求方式为POST
    if flask.request.method == "POST":
        # 判断表单中是否上传了名为image的文件
        if flask.request.files.get("image"):
            # 1. 读取上传图片二进制流
            image_bytes = flask.request.files["image"].read()
            # 二进制流转PIL图像对象
            image = Image.open(io.BytesIO(image_bytes))
            # 2. 图像预处理到模型输入格式
            img_input = prepare_image(image, target_size=(224, 224))

            # 3. 模型推理，关闭梯度计算节省显存
            with torch.no_grad():
                pred_output = model(img_input)
            # softmax转0~1概率分布
            preds = F.softmax(pred_output, dim=1)
            # 取概率Top3结果
            top3_vals, top3_idx = torch.topk(preds.cpu().data, k=3, dim=1)
            # 转numpy脱离计算图
            top3_vals = top3_vals.detach().cpu().numpy()
            top3_idx = top3_idx.detach().cpu().numpy()

            # 4. 组装预测结果列表
            data["predictions"] = []
            for prob, label in zip(top3_vals[0], top3_idx[0]):
                result_item = {
                    "label": str(label),
                    "probability": float(prob)
                }
                data["predictions"].append(result_item)
            # 标记预测成功
            data["success"] = True

    # 统一封装标准JSON返回体
    resp_json = {
        "message": "success",
        "code": 200,
        "data": data
    }
    return flask.jsonify(resp_json)


# GET测试接口
@app.route(rule="/get_demo", methods=["GET"])
def get_demo():
    name = flask.request.args.get("name", "匿名用户")
    age = flask.request.args.get("age", "未知")
    res = {
        "method": "GET",
        "msg": "这是GET测试接口，参数拼在URL里",
        "你的姓名": name,
        "你的年龄": age
    }
    return flask.jsonify(res)

# JSON POST测试接口
@app.route(rule="/post_demo", methods=["POST"])
def post_demo():
    json_data = flask.request.get_json()
    if not json_data:
        return flask.jsonify({"success": False, "msg": "请传入JSON参数"})
    username = json_data.get("username")
    text = json_data.get("text")
    res = {
        "method": "POST",
        "msg": "这是普通JSON POST接口，数据藏在请求体中",
        "接收用户名": username,
        "接收文本内容": text,
        "success": True
    }
    return flask.jsonify(res)


# 程序入口
if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    load_model()  # 启动前预加载模型，避免每次请求重复加载
    # host=0.0.0.0 允许局域网其他设备访问；port=5012指定端口；debug开发调试模式
    app.run(host='0.0.0.0', port=5012, debug=True)

客户端调用代码（flask_server.py）

python 复制代码 
import requests
# 本地测试用127.0.0.1，同局域网其他电脑访问改为局域网IP 192.168.31.33
flask_url = 'http://127.0.0.1:5012/predict'

def predict_result(image_path):
    try:
        # rb二进制只读打开图片
        image = open(image_path, 'rb').read()
        # files参数专门传文件，对应form-data
        payload = {'image': image}
        r = requests.post(flask_url, files=payload).json()

        if r['data']['success']:
            # 循环打印top3预测结果
            for i, result in enumerate(r["data"]['predictions']):
                print('{}.预测类别为{}，概率：{}'.format(i + 1, result['label'], result['probability']))
        else:
            print('接口推理失败')
    except Exception as e:
        print("请求异常：", e)

if __name__ == '__main__':
    # 替换成你的图片路径
    predict_result('./flower_data/val_filelist/image_00028.jpg')

独立预测脚本（pyyy.py）------用于部署前验证模型加载是否正确

python 复制代码 
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
from torchvision import transforms, models

model = None
use_gpu = False

# 加载预训练ResNet18模型+权重
def load_model():
    """Load the pre-trained model, you can use your model just as easily."""
    global model
    # 1. 初始化resnet18主干网络
    model = models.resnet18()
    # 2. 获取原模型最后全连接层输入维度
    num_ftrs = model.fc.in_features
    # 3. 替换全连接层，适配102分类任务
    model.fc = nn.Sequential(nn.Linear(num_ftrs, 102))
    # 4. 读取训练好的最优权重文件best.pth
    checkpoint = torch.load('best.pth', weights_only=True)
    # 5. 将权重加载到模型中
    model.load_state_dict(checkpoint['state_dict'])
    # 6. 设置模型为评估模式（关闭dropout、batchnorm训练行为）
    model.eval()
    # 7. 如果开启GPU，模型移至cuda
    if use_gpu:
        model.cuda()


# 图像预处理函数
def prepare_image(image, target_size):
    # 统一转为RGB三通道图像
    if image.mode != 'RGB':
        image = image.convert('RGB')
    # 1. 缩放到模型输入尺寸224*224
    image = transforms.Resize(target_size)(image)
    # 2. PIL图像转张量，像素值从0~255映射到0~1
    image = transforms.ToTensor()(image)
    # 3. ImageNet标准化，和训练时保持一致
    image = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image)
    # 4. 增加batch维度 [C,H,W] -> [1,C,H,W]，适配模型输入格式
    image = image[None]
    # GPU数据迁移
    if use_gpu:
        image = image.cuda()
    return image


# 预测主函数，输入PIL图像，返回json格式预测结果
def predict(image):
    # 初始化返回字典，success标记是否预测成功
    data = {"success": False}
    # 图像预处理
    image = prepare_image(image, target_size=(224, 224))
    # 前向推理，softmax把logits转为各类别概率
    preds = F.softmax(model(image), dim=1)
    # 取概率前3大的类别，返回(概率值张量,类别索引张量)
    results = torch.topk(preds.cpu(), k=3, dim=1)
    # 分离计算图、移到CPU、转numpy数组，解除梯度占用内存
    results = (results[0].detach().cpu().numpy(), results[1].detach().cpu().numpy())
    # 存放top3预测结果
    data['predictions'] = list()
    # 遍历top3概率+对应标签
    for prob, label in zip(results[0][0], results[1][0]):
        r = {"label": str(label), "probability": float(prob)}
        data['predictions'].append(r)
    # 修改成功标记
    data["success"] = True
    return data


# 程序入口：单独运行该脚本时执行测试代码
if __name__ == '__main__':
    # 第一步加载模型权重
    load_model()
    # 测试图片路径
    img_path = r"D:\Jupyter_Projects\PythonProject\模型部署\flower_data\train_filelist\image_00001.jpg"
    # 打开PIL图像
    test_img = Image.open(img_path)
    # 执行预测
    res = predict(test_img)
    # 打印预测结果
    print(res)

二、依赖库导入与功能分层

python 复制代码 
import io
import flask
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
from torchvision import transforms, models

逐库功能解析：

  • flask ：轻量级Web框架，用于构建HTTP服务。flask.Flask初始化应用对象，flask.request获取请求数据，flask.jsonify将Python字典序列化为JSON响应。

  • torch / torch.nn.functional ：PyTorch核心库。torch.no_grad()上下文管理器在推理时关闭梯度计算，大幅减少显存占用并加速前向传播；F.softmax将模型输出的logits转换为0~1之间的概率分布。

  • PIL（Pillow） ：图像处理库。Image.open将二进制流解码为PIL图像对象，便于后续transforms流水线处理。

  • torchvision.transforms / models ： torchvision官方计算机视觉工具库。transforms.Compose构建图像预处理流水线（缩放、张量化、标准化）；models.resnet18加载ResNet18网络结构。

三、模型加载与全局配置

python 复制代码 
app = flask.Flask(__name__)
model = None
use_gpu = False
idx2label = {}

def load_model():
    global model, use_gpu
    model = models.resnet18(weights=None)
    model.eval()
    if torch.cuda.is_available():
        use_gpu = True
        model.cuda()

核心要点解析：

  1. 全局变量设计 ：model和use_gpu声明为全局变量，在load_model()中赋值，后续预测接口中直接使用，避免每次请求重复加载模型。

  2. weights=None的用意 ：表示不加载ImageNet预训练权重，因为我们会用自己的best.pth覆盖。如果使用迁移学习，可改为weights=models.ResNet18_Weights.IMAGENET1K_V1加载官方预训练权重。

  3. model.eval() ：将模型切换为评估模式，关键作用是关闭Dropout和BatchNorm的训练行为（如Dropout不再随机失活，BatchNorm使用全局统计量而非当前batch统计量），保证推理结果的确定性。

  4. GPU检测与迁移 ：torch.cuda.is_available()检测当前环境是否有可用GPU，若有则将模型移至GPU。实际生产环境中，若GPU显存紧张，可强制使用CPU推理。

四、图像预处理流水线

python 复制代码 
def prepare_image(image, target_size):
    transform = transforms.Compose([
        transforms.Resize(target_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    img_tensor = transform(image)
    img_tensor = torch.unsqueeze(img_tensor, 0)
    if use_gpu:
        img_tensor = img_tensor.cuda()
    return img_tensor

预处理三步走：

  1. Resize ：将所有输入图片统一缩放到224×224像素，这是ResNet18的标准输入尺寸。

  2. ToTensor ：将PIL图像（像素值0255）转换为PyTorch张量（像素值01），并自动将形状从(H, W, C)调整为(C, H, W)。

  3. Normalize ：使用ImageNet数据集的均值和标准差进行标准化。这一步至关重要------即便我们是在花卉数据集上微调的模型，训练时同样使用了这套标准化参数，推理时必须保持一致，否则模型看到的输入分布与训练时不同，预测结果会严重偏差。

  4. 增加batch维度 ：torch.unsqueeze(img_tensor, 0)将形状从[C, H, W]变为[1, C, H, W]，因为PyTorch模型要求输入为batch格式。

五、核心预测接口深度解析

python 复制代码 
@app.route("/predict", methods=["POST"])
def predict():
    data = {"success": False}
    if flask.request.method == "POST":
        if flask.request.files.get("image"):
            image_bytes = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image_bytes))
            img_input = prepare_image(image, target_size=(224, 224))
            
            with torch.no_grad():
                pred_output = model(img_input)
            preds = F.softmax(pred_output, dim=1)
            top3_vals, top3_idx = torch.topk(preds.cpu().data, k=3, dim=1)
            top3_vals = top3_vals.detach().cpu().numpy()
            top3_idx = top3_idx.detach().cpu().numpy()
            
            data["predictions"] = []
            for prob, label in zip(top3_vals[0], top3_idx[0]):
                data["predictions"].append({
                    "label": str(label),
                    "probability": float(prob)
                })
            data["success"] = True
    
    return flask.jsonify({
        "message": "success",
        "code": 200,
        "data": data
    })

接口设计逐层拆解：

  1. 路由与请求方法限制 ：@app.route("/predict", methods=["POST"])限定该接口仅接受POST请求，GET请求会返回405 Method Not Allowed。

  2. 文件读取 ：flask.request.files.get("image")从multipart/form-data表单中获取名为image的文件对象。注意：客户端必须使用form-data格式上传，而非JSON格式。

  3. 二进制流转PIL图像 ：Image.open(io.BytesIO(image_bytes))将二进制数据解码为PIL图像对象。

  4. 推理上下文 ：with torch.no_grad():包裹前向传播代码，在此上下文中的所有操作都不会构建计算图，既不保存梯度信息，显著降低显存消耗。

  5. Top-K提取 ：torch.topk(preds.cpu().data, k=3, dim=1)在类别维度（dim=1）上取出概率最大的前3个类别。

  6. 张量转Python原生类型 ：detach()切断计算图引用→cpu()移至CPU→numpy()转为numpy数组→float()转为Python float。这一系列操作确保数据可被jsonify正常序列化。

  7. 统一返回体 ：所有接口返回格式保持一致------{ "message": "...", "code": 200, "data": {...} }，便于前端统一处理。

六、辅助测试接口：GET与POST JSON

python 复制代码 
@app.route(rule="/get_demo", methods=["GET"])
def get_demo():
    name = flask.request.args.get("name", "匿名用户")
    age = flask.request.args.get("age", "未知")
    return flask.jsonify({
        "method": "GET",
        "msg": "这是GET测试接口，参数拼在URL里",
        "你的姓名": name,
        "你的年龄": age
    })

@app.route(rule="/post_demo", methods=["POST"])
def post_demo():
    json_data = flask.request.get_json()
    if not json_data:
        return flask.jsonify({"success": False, "msg": "请传入JSON参数"})
    username = json_data.get("username")
    text = json_data.get("text")
    return flask.jsonify({
        "method": "POST",
        "msg": "这是普通JSON POST接口，数据藏在请求体中",
        "接收用户名": username,
        "接收文本内容": text,
        "success": True
    })

两个测试接口的设计目的：

  • GET接口 ：演示如何从URL查询参数中取值（flask.request.args.get()），用于快速验证服务是否正常运行、网络是否通畅。

  • POST JSON接口 ：演示如何解析请求体中的JSON数据（flask.request.get_json()），用于测试JSON格式的POST请求是否被正确解析。

这两个接口在开发阶段非常实用------当/predict接口出现问题时，可以先通过这两个接口排除Flask服务本身的故障。

七、服务启动与运行

python 复制代码 
if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    load_model()
    app.run(host='0.0.0.0', port=5012, debug=True)

启动参数详解：

  • host='0.0.0.0' ：监听所有网络接口，不仅限127.0.0.1。同局域网内的其他设备（如手机、另一台电脑）可通过http://<本机IP>:5012访问服务。

  • port=5012：指定服务端口为5012，避免与常用端口（如80、8080）冲突。

  • debug=True ：开启调试模式，代码修改后服务会自动重载，且错误页面会显示详细堆栈信息，极大方便开发调试。生产环境务必设置为False，否则会带来安全风险。

八、客户端调用方式详解

1. Python脚本调用

python 复制代码 
import requests

flask_url = 'http://127.0.0.1:5012/predict'

def predict_result(image_path):
    image = open(image_path, 'rb').read()
    payload = {'image': image}
    r = requests.post(flask_url, files=payload).json()
    if r['data']['success']:
        for i, result in enumerate(r["data"]['predictions']):
            print(f'{i+1}.预测类别为{result["label"]}，概率：{result["probability"]}')

关键点：

  • files=payload会自动将数据编码为multipart/form-data格式，与服务端的request.files对应。
  • 返回的r是Python字典，可直接通过键名访问嵌套数据。

2. curl命令测试

测试POST JSON接口：

bash 复制代码 
curl -X POST -H "Content-Type:application/json" -d "{\"username\":\"学生\",\"text\":\"success\"}" http://127.0.0.1:5012/post_demo

测试GET接口：

bash 复制代码 
curl "http://127.0.0.1:5012/get_demo?name=123&age=111"

测试预测接口（上传图片）：

bash 复制代码 
curl -X POST -F "image=@./flower_data/val_filelist/image_00028.jpg" http://127.0.0.1:5012/predict

运行结果：

3. Postman图形化调试

  1. 选择请求方法为POST，输入URLhttp://127.0.0.1:5012/predict。
  2. Body选择form-data，Key填image，Value选择文件（图片）。
  3. 点击Send，即可在下方看到JSON响应。

九、调试结果分析与常见问题排查

1. 预测结果解读

从实际运行结果来看：

json 复制代码 
{
    "message": "success",
    "code": 200,
    "data": {
        "success": true,
        "predictions": [
            {"label": "673", "probability": 0.0385},
            {"label": "348", "probability": 0.0380},
            {"label": "xxx", "probability": 0.0375}
        ]
    }
}

  • 概率值普遍偏低（如0.038） ：说明模型对该图片的置信度不高，这是模型质量问题，接口逻辑本身是正常的。

  • label为数字索引 ：因为代码中未补全idx2label映射字典。若想显示可读的类别名称，可在服务端维护该字典（从训练集获取），并在返回时将str(label)替换为idx2label[int(label)]。

2. 文件上传大小限制

Flask默认限制上传文件大小为16MB，若图片过大，可在app.run之前设置：

python 复制代码 
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024  # 100MB

3. 模型加载路径错误

torch.load('best.pth')使用相对路径，启动服务时需确保当前工作目录包含该文件，否则报FileNotFoundError。建议使用绝对路径或通过环境变量配置模型路径。

4. GPU显存不足

如果服务同时接收多个并发请求，GPU显存可能耗尽。可在predict函数中每次推理后手动清理：

python 复制代码 
if use_gpu:
    torch.cuda.empty_cache()

但频繁调用empty_cache()会影响性能，更推荐的做法是限制并发数或使用异步推理队列。

5. 跨域问题（CORS）

如果前端是分离的Web应用（如Vue、React），直接调用Flask接口会因跨域策略被浏览器拦截。可使用flask-cors库解决：

bash 复制代码 
pip install flask-cors
python 复制代码 
from flask_cors import CORS
CORS(app)  # 允许所有跨域请求

十、项目文件结构

复制代码 
project/
├── flask_predict.py          # 服务端主程序
├── flask_server.py           # 客户端调用示例
├── best.pth                  # 训练好的模型权重（102分类）
├── flower_data/              # 测试图片目录
│   ├── train_filelist/
│   │   └── image_00001.jpg
│   └── val_filelist/
│       └── image_00028.jpg
└── pyyy.py                   # 独立预测脚本（验证模型加载用）

其中pyyy.py是独立于Flask的预测脚本，用于在部署前验证模型加载和预处理是否正确，不依赖Web环境，适合单独调试模型本身。

总结

本文完整展示了基于Flask部署PyTorch图像分类模型的全部流程，涵盖：

  • 模型加载：如何加载训练好的权重文件，适配GPU/CPU环境。
  • 图像预处理：标准化、尺寸调整、张量转换的完整流水线。
  • 接口设计 ：核心预测接口/predict的完整实现，包括文件读取、推理、Top-K结果提取、JSON序列化。
  • 辅助调试：GET/POST测试接口的设计思路与使用方法。
  • 多种调用方式：Python脚本、curl命令、Postman图形化工具。
  • 问题排查：文件大小限制、模型路径、GPU显存、跨域等常见问题的解决方案。

通过以上实践，读者可以快速掌握模型服务化的基本技能，并将此套代码模板灵活迁移到自己的项目中。任何部署问题，都可以先通过/get_demo和/post_demo测试接口逐步排查，确保Flask服务本身正常后，再定位模型推理环节的问题。

相关推荐
我的温馨家园1 小时前
新手入局干细胞研究的五大认知陷阱
大数据·人工智能·精选
爱吃大芒果1 小时前
AI 智能体工作流设计蓝图：将非结构化情绪记录转化为高精度的模型 Prompt 上下文
人工智能·华为·prompt·harmonyos
小尘要自信1 小时前
小米摄像头怎么接入RTSP？Go2RTC转流、EasyNVR录像与公网访问教程
人工智能·docker·开源
ACP广源盛139246256731 小时前
GSV6155 @ACP#工业车规 DP1.4 重定时器 Retimer
大数据·人工智能·分布式·嵌入式硬件
Black蜡笔小新1 小时前
企业AI算力工作站/企业级AI模型工作站DLTM训推一体工作站破解企业AI建模难题
人工智能·机器学习
海兰1 小时前
【AI编程思考：第八篇】从演示到生产：AI工程部署实战指南
大数据·人工智能·ai编程
杨超越luckly1 小时前
Agent 应用指南：基于 OurAirports 的中国机场设施数据可视化
python·html·github·可视化·机场设施
卷无止境1 小时前
pytest 从零到实战：要想代码好，测试少不了
后端·python
水龙吟啸1 小时前
华为2026.7.1机考选择题+编程题【速刷敲黑板】
人工智能
热门推荐
012026 年 AI 编程工具终极横评：Cursor vs Claude Code vs Copilot vs Windsurf02GitHub 镜像站点03AI科技热点日报 | 2026年07月01日04【AI】2026 年具身智能模型和世界模型总结052026 年 AI 大模型 & AI 编程工具实战全总结062026 AI 编程工具选型横评：Cursor / Claude Code / Trae / Copilot 到底选谁（建议收藏·避坑版）07AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路082026 AI 编程工具终极实战指南：Cursor vs Claude Code vs Copilot，开发者该怎么选？092026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？102026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片