用CodeBuddy开发 AI模型部署脚本：TensorFlow模型转ONNX并部署到Python服务

友友们，早上好啊。今天继续AI工具实战系列。我将分享最近几个月使用ai工具用于工作的那些事。既是一次知识的分享，又是一次自我的一次总结，也希望自己的一些使用经验可以帮助到大家。下面正文开始

快要放国庆了，工作兴趣直线下降。是时候用ai继续偷懒了，本篇文章我主要想做的就是用ai工具来服务ai，说干就干。

一、提出目标

场景

这次的使用场景是在我们航道系统中船舶图片智能分类子系统中，原TensorFlow图像分类模型（.h5格式）因推理速度不足（单张图片处理需1200ms）导致用户体验下降。我想着把这个性能提升一下，于是决定将模型转换为ONNX格式（预期推理速度提升30%），并开发Python Flask服务提供标准化推理接口。

先列出具体目标：

1. 格式转换：将.h5 格式模型转为 ONNX 格式，确保输入输出维度、数据类型与原模型完全一致，避免推理逻辑偏差；
2. 服务开发：基于 Python Flask 构建推理服务，提供 POST 类型的 /predict 接口，支持接收图片 Base64 编码作为输入；
3. 结果输出：接口需返回结构化 JSON 数据，包含分类类别（如 "小型船"）与置信度（如 0.98），且推理响应时间控制在 1 秒内。

二、使用工具

同样我们还是来看下技术栈和工具：

• 模型处理：Python 3.10（基础环境）、TensorFlow 2.15（加载.h5 模型）、tf2onnx（格式转换）、onnxruntime（ONNX 模型推理）；
• 服务开发：Flask（轻量级 Web 框架，快速构建 API）；
• 协作工具：CodeBuddy（AI 辅助代码生成与问题排查）。

三、协作关键步骤

1. 需求输入与代码生成

按照我们之前的经验，AI 工具的协作效率，始于精准的需求描述。于是我们需要向 CodeBuddy 输入结构化需求，明确每个环节的技术细节与约束条件： “请协助开发两套 Python 脚本，具体要求如下：

1. 模型转换脚本（model_convert.py）：使用 tf2onnx 库加载 TensorFlow .h5 模型，转换为 ONNX 格式。需指定输入形状为 (1,224,224,3)（批量大小 1、图像尺寸 224x224、RGB 三通道），确保转换后模型的输入输出格式与原模型一致；
2. Flask 服务脚本（flask_service.py）：开发 POST 接口 /predict，流程包含四步：接收请求中的图片 Base64 编码→解码为 PIL 图像→预处理（resize 至 224x224、转 RGB 通道、归一化处理、调整维度为 (1,224,224,3)）→调用 onnxruntime 执行推理→返回 {"class": "类别名称", "confidence": 置信度} 格式结果；
3. 所有代码需兼容 Python 3.10 + TensorFlow 2.15，注释清晰，便于后续调试。”

提出明确需求

CodeBuddy 在 很快生成了两套脚本，核心逻辑如下：

• model_convert.py：通过tf.keras.models.load_model()加载.h5 模型，调用tf2onnx.convert.from_keras()指定输入形状，生成 ONNX 文件；
• flask_service.py：基于 Flask 构建服务，使用base64.b64decode()解码图片，PIL.Image处理图像，onnxruntime.InferenceSession()执行推理。

2 解决版本兼容问题

但运行 model_convert.py 时，终端立即抛出错误：AttributeError: module 'tensorflow' has no attribute 'compat'。将错误信息直接反馈给 CodeBuddy 后，其迅速定位问题根源 ——tf2onnx 版本与 TensorFlow 2.15 不兼容，并给出解决方案：“安装 tf2onnx==1.14.0 版本，该版本针对 TensorFlow 2.15 做了适配，可避免 compat 模块调用错误”。

报错修改

执行pip install tf2onnx==1.14.0后重新运行脚本，成功生成 ONNX 模型文件，且通过onnxruntime.get_inputs()验证，输入维度与数据类型均与原模型一致。

3 调试推理结果异常

ONNX 模型生成后，启动 Flask 服务进行接口测试：通过 Postman 发送 Base64 编码的猫类图片，返回结果却显示{"class": "dog", "confidence": 0.0}，置信度异常为 0。

向 CodeBuddy 提问：“ONNX 模型推理时置信度始终为 0，可能的原因有哪些？”

继续提问

其立即给出排查方向：“优先检查图像预处理步骤与模型训练时是否一致，尤其是归一化数值范围（如训练时用 0-1 范围，推理时却用 0-255，会导致输入数据分布偏差）”。

查看 Flask 脚本的预处理代码，发现image = image / 255.0的归一化逻辑，而原模型训练时使用的是tf.keras.applications.resnet50.preprocess_input()（将像素值映射到 - 1 至 1 范围）。修正代码为image = tf.keras.applications.resnet50.preprocess_input(image)后重新测试，接口成功返回{"class": "cat", "confidence": 0.98}，结果准确率与原 TensorFlow 模型完全一致。

改进后代码：

模型转换脚本（model_convert.py）

import tensorflow as tf
from tf2onnx import convert

def convert_h5_to_onnx(h5_model_path, onnx_output_path):
    # 加载TensorFlow .h5模型
    model = tf.keras.models.load_model(h5_model_path)
    print(f"成功加载.h5模型：{h5_model_path}")
    
    # 转换为ONNX格式，指定输入形状（1,224,224,3）
    convert.from_keras(
        model,
        input_signature=[tf.TensorSpec(shape=(1, 224, 224, 3), dtype=tf.float32)],
        output_path=onnx_output_path,
        opset=16  # 适配onnxruntime的opset版本
    )
    print(f"ONNX模型已保存至：{onnx_output_path}")

if __name__ == "__main__":
    # 替换为实际的模型路径
    H5_MODEL_PATH = "image_classifier.h5"
    ONNX_OUTPUT_PATH = "image_classifier.onnx"
    convert_h5_to_onnx(H5_MODEL_PATH, ONNX_OUTPUT_PATH)

可以看到，该脚本的关键在于input_signature参数的设置，通过tf.TensorSpec明确输入形状与数据类型，确保 ONNX 模型的输入接口与原模型完全匹配，避免后续推理时的维度错误。

Flask 推理服务脚本（flask_service.py）

from flask import Flask, request, jsonify
import base64
from io import BytesIO
from PIL import Image
import onnxruntime as ort
import numpy as np
import tensorflow as tf

app = Flask(__name__)

# 加载ONNX模型，初始化推理会话
ONNX_MODEL_PATH = "image_classifier.onnx"
session = ort.InferenceSession(ONNX_MODEL_PATH)
# 获取模型输入名称（适配不同模型的输入命名差异）
input_name = session.get_inputs()[0].name
# 类别映射（需与训练时的类别顺序一致）
CLASS_MAPPING = {0: "cat", 1: "dog", 2: "bird"}

def preprocess_image(base64_str):
    """将Base64编码转换为模型所需的输入格式"""
    # 1. Base64解码
    img_data = base64.b64decode(base64_str)
    # 2. 转为PIL图像并调整通道为RGB
    img = Image.open(BytesIO(img_data)).convert("RGB")
    # 3. Resize至224x224
    img = img.resize((224, 224))
    # 4. 转为numpy数组并调整维度（(224,224,3) → (1,224,224,3)）
    img_array = np.array(img, dtype=np.float32)
    img_array = np.expand_dims(img_array, axis=0)
    # 5. 归一化（与训练时保持一致：使用ResNet50的预处理逻辑）
    img_array = tf.keras.applications.resnet50.preprocess_input(img_array)
    return img_array

@app.route("/predict", methods=["POST"])
def predict():
    # 1. 接收请求数据
    data = request.json
    if "image_base64" not in data:
        return jsonify({"error": "请求缺少image_base64字段"}), 400
    
    # 2. 图像预处理
    try:
        input_data = preprocess_image(data["image_base64"])
    except Exception as e:
        return jsonify({"error": f"图像预处理失败：{str(e)}"}), 400
    
    # 3. ONNX模型推理
    try:
        outputs = session.run(None, {input_name: input_data})
        # 4. 结果后处理（获取置信度最高的类别）
        pred_probs = outputs[0][0]  # 输出形状：(1, 类别数) → 取第一个样本的概率
        pred_class_idx = np.argmax(pred_probs)
        pred_class = CLASS_MAPPING[pred_class_idx]
        pred_confidence = round(float(pred_probs[pred_class_idx]), 4)
    except Exception as e:
        return jsonify({"error": f"模型推理失败：{str(e)}"}), 500
    
    # 5. 返回结构化结果
    return jsonify({
        "class": pred_class,
        "confidence": pred_confidence
    })

if __name__ == "__main__":
    # 调试模式：host=0.0.0.0允许外部访问，port=5000为默认端口
    app.run(host="0.0.0.0", port=5000, debug=True)

一起来看一下，我们可以发现，脚本中两个关键设计：一是通过session.get_inputs()[0].name动态获取输入名称，避免因模型命名差异导致的推理错误；二是将预处理逻辑封装为独立函数，便于后续修改与维护，同时明确标注与训练时的对齐点（如 ResNet50 的预处理逻辑）。

四、最终效果验证

格式验证：使用onnx.checker.check_model()检查 ONNX 模型完整性，确认无格式错误；

速度对比：在相同硬件环境（CPU：Intel i7-12700H，内存：16GB）下，原 TensorFlow 模型推理平均耗时 1.2 秒，ONNX 模型平均耗时 0.84 秒，推理速度提升 30%；

接口测试：通过 Postman 发送 Base64 编码的测试图片，1 秒内收到返回结果，且分类准确率与原模型完全一致（如小型船类图片返回{"class": "小型船", "confidence": 0.98}）。

效果对比

五、经验总结

这里我主要总结一下，本次遇到的推理结果异常问题，这是部署的常见问题。遇到这类问题，我们首先要检查预处理一致性：对比推理时与训练时的图像尺寸、通道顺序、归一化范围，确保输入数据分布一致；然后验证模型输入输出：使用onnxruntime.get_inputs()与tf.keras.Model.summary()对比原模型与 ONNX 模型的输入维度、数据类型；最后调试中间结果：在代码中打印预处理后的图像数组、模型推理的中间输出，定位数据异常的具体环节。

文末还是说下AI工具吧。目前来说AI工具仍需人工介入，如类别映射（CLASS_MAPPING）需根据实际训练数据配置，预处理逻辑需结合训练过程调整，这些 业务特异性环节无法完全依赖 AI 完成。好了，本期的分享就先到这里把，我们下期见！！