如何提取tflite模型权重并转换为C变量供STM32 X-Cube-AI使用

你需要的功能可以通过解析TFLite模型结构、单独提取常量权重段实现，以下是完整实现方案：

核心思路 #

• 调用TensorFlow Lite官方Python接口解析模型，过滤输入输出等临时张量，仅提取网络的权重、偏置等固定常量参数
• 按照STM32 Flash 4字节对齐的访问要求，对拼接后的权重数据做对齐补全，匹配X-Cube AI的存储格式
• 直接输出符合C语法的常量数组，可直接加入STM32工程编译，或通过UART传输后写入对应Flash地址

Python实现脚本 #

import tensorflow as tf
import numpy as np

def extract_tflite_weights_to_c(tflite_path, output_c_path, var_name="ai_model_weights"):
    # 加载TFLite模型并初始化解析器
    interpreter = tf.lite.Interpreter(model_path=tflite_path)
    interpreter.allocate_tensors()

    weight_bytes_list = []
    # 遍历所有张量，筛选权重类常量
    for tensor_info in interpreter.get_tensor_details():
        # 跳过输入、输出等激活张量，仅保留固定常量参数
        if tensor_info['name'] and ('weight' in tensor_info['name'].lower() or 'bias' in tensor_info['name'].lower()):
            tensor_data = interpreter.get_tensor(tensor_info['index'])
            weight_bytes_list.append(tensor_data.tobytes())

    # 拼接所有权重并做4字节对齐，适配STM32存储要求
    total_weight_bytes = b''.join(weight_bytes_list)
    pad_len = (4 - len(total_weight_bytes) % 4) % 4
    total_weight_bytes += b'\x00' * pad_len

    # 生成C文件内容
    c_file_content = f"""/* Auto-generated TFLite weights for STM32 X-Cube AI */
#include <stdint.h>

/* 网络权重数组，4字节对齐 */
const uint8_t {var_name}[{len(total_weight_bytes)}] __attribute__((aligned(4))) = {{
"""
    # 按16字节每行格式化输出
    for idx in range(0, len(total_weight_bytes), 16):
        chunk = total_weight_bytes[idx:idx+16]
        line = "    " + ", ".join([f"0x{byte_val:02X}" for byte_val in chunk]) + ",\n"
        c_file_content += line

    c_file_content += "};\n"
    # 追加权重总长度变量，方便UART传输时读取
    c_file_content += f"const uint32_t {var_name}_total_len = {len(total_weight_bytes)};\n"

    # 写入输出文件
    with open(output_c_path, "w", encoding="utf-8") as f:
        f.write(c_file_content)

# 使用示例
if __name__ == "__main__":
    # 替换为你的tflite模型路径和输出C文件路径
    extract_tflite_weights_to_c("./test_model.tflite", "./ai_weights.c")

适配说明 #

• 脚本默认按权重、偏置的名称筛选张量，如果你需要和X-Cube AI自动生成的权重顺序完全一致，可以先导出一次X-Cube AI示例工程，对比工程中network_weights.c的内容调整张量的筛选和排序规则
• 生成的数组默认带const修饰，可直接存储到STM32内部Flash；如果需要加载到RAM运行，删除const修饰符即可
• 输出的字节序为小端，和STM32的运行字节序完全匹配，不需要额外做字节序转换
• 如果你需要通过UART传输权重，直接读取total_weight_bytes变量的二进制内容发送即可，不需要转成C格式

内容的提问来源于stack exchange，提问作者littleLad