3. 卷积神经网络权重提取与量化

任务目标

从训练好的Keras CNN模型中提取卷积层与全连接层的权重和偏置，将其量化为int8格式，并生成可直接嵌入MCU工程的C头文件。本案例以MNIST手写数字识别模型为例，演示完整的提取-量化-导出流程。

环境依赖

tensorflow==2.13.0
numpy==1.24.3

实现代码

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models

# 1. 构建并训练一个简单CNN（仅用于演示，实际使用请加载自己的模型）
def build_and_train_mnist_model():
    model = models.Sequential([
        layers.Conv2D(8, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        layers.MaxPooling2D((2, 2)),
        layers.Flatten(),
        layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255.0
    model.fit(x_train[:2000], y_train[:2000], epochs=1, verbose=0)
    return model

# 2. 量化函数：将浮点权重映射到int8范围[-128, 127]
def quantize_tensor(tensor):
    min_val = np.min(tensor)
    max_val = np.max(tensor)
    scale = (max_val - min_val) / 255.0
    zero_point = -128 - round(min_val / scale) if scale > 0 else 0
    quantized = np.round(tensor / scale + zero_point).astype(np.int8)
    return quantized, scale, zero_point

# 3. 提取并量化所有层参数
def extract_and_quantize(model):
    layer_params = []
    for layer in model.layers:
        weights = layer.get_weights()
        if len(weights) == 0:
            continue
        layer_info = {'name': layer.name, 'type': layer.__class__.__name__}
        quantized_weights = []
        for w in weights:
            q, scale, zp = quantize_tensor(w)
            quantized_weights.append({'data': q, 'scale': scale, 'zero_point': zp})
        layer_info['params'] = quantized_weights
        layer_params.append(layer_info)
    return layer_params

# 4. 生成C头文件
def generate_c_header(layer_params, filename="mnist_cnn_weights.h"):
    header = "#ifndef MNIST_CNN_WEIGHTS_H\n#define MNIST_CNN_WEIGHTS_H\n\n"
    header += "#include <stdint.h>\n\n"
    
    for idx, layer in enumerate(layer_params):
        for p_idx, param in enumerate(layer['params']):
            data = param['data'].flatten()
            var_name = f"{layer['name']}_{p_idx}"
            header += f"// {layer['type']} - scale: {param['scale']:.8f}, zero_point: {param['zero_point']}\n"
            header += f"const int8_t {var_name}[{len(data)}] = {{"
            header += ", ".join(str(v) for v in data)
            header += "};\n\n"
    
    header += "#endif // MNIST_CNN_WEIGHTS_H\n"
    
    with open(filename, 'w') as f:
        f.write(header)
    print(f"Header file generated: {filename}")

# 主流程
if __name__ == "__main__":
    # 实际使用时替换为：model = tf.keras.models.load_model('your_model.h5')
    model = build_and_train_mnist_model()
    layer_params = extract_and_quantize(model)
    generate_c_header(layer_params)
    
    # 验证：打印各层量化参数统计
    for layer in layer_params:
        for p_idx, param in enumerate(layer['params']):
            data = param['data']
            print(f"{layer['name']}_{p_idx}: shape={data.shape}, "
                  f"min={data.min()}, max={data.max()}, "
                  f"scale={param['scale']:.6f}, zp={param['zero_point']}")

运行与验证

运行方式：

python extract_quantize_weights.py

验证方法：

1. 程序正常结束并生成mnist_cnn_weights.h文件；
2. 控制台输出每层参数的shape、量化后的int8取值范围（应在-128~127之间）、scale和zero_point值；
3. 打开生成的C头文件，确认包含#include <stdint.h>、各层权重数组声明及注释中的scale/zero_point信息；
4. 将头文件中的数组与原始浮点权重进行反量化对比：float_val = (int8_val - zero_point) * scale，误差应在±1/scale范围内（即量化精度损失在预期内）。

关键说明

• 量化采用对称线性量化，将浮点范围映射到[-128, 127]；实际部署时需在MCU端保存每个张量的scale和zero_point用于反量化；
• 本案例生成的C头文件仅包含权重数据，不包含推理代码；MCU端需自行实现卷积和全连接运算，并按上述公式反量化；
• 若模型包含BatchNormalization层，建议在训练后将BN层折叠进前一卷积层再提取，否则需额外处理；
• 对于更大模型，建议将权重存储为.c文件而非头文件，避免编译时内存占用过高。