第4章：TensorFlow模型设计与训练

好，咱们进入正题。这一章要聊的是怎么在Google Colab里搭一个简单的MLP模型，然后训练、评估，最后转成.tflite格式。说白了，就是让我们的AI模型能在STM32上跑起来。

4.1 为什么选Google Colab？

我个人习惯用Colab做模型训练。原因很简单——免费GPU。你想想看，自己配一台带GPU的电脑，少说也得大几千。Colab直接给你白嫖，何乐而不为？

我在项目中遇到过不少次，本地训练模型跑到一半，电脑直接卡死。换了Colab之后，再也没这烦恼了。而且它跟Google Drive无缝对接，模型文件直接存云端，方便得很。

4.2 搭建简单MLP模型

MLP，全称多层感知机。说白了就是一堆神经元堆叠起来。咱们这次要做的，是一个三层的MLP。

代码长这样：

import tensorflow as tf
from tensorflow import keras

model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(input_dim,)),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dense(num_classes, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

嗯，这里要注意。input_dim和num_classes要根据你的实际数据来定。比如你用了加速度计、陀螺仪、磁力计三个传感器，每个输出3轴数据，那input_dim就是9。

4.3 训练数据准备

数据准备这一步，我吃过不少亏。曾经有一次，我偷懒没做数据归一化，结果模型训练了三天，准确率死活上不去。后来才发现，数据范围差太多，模型根本学不动。

数据准备的步骤：

1. 从STM32采集原始数据，存成CSV文件
2. 上传到Google Drive
3. 在Colab里挂载Drive，读取数据
4. 做归一化处理
5. 划分训练集和测试集

代码示例：

from google.colab import drive
drive.mount('/content/drive')

import pandas as pd
import numpy as np

# 读取数据
data = pd.read_csv('/content/drive/MyDrive/sensor_data.csv')

# 归一化
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
X = scaler.fit_transform(data.iloc[:, :-1])
y = data.iloc[:, -1].values

# 划分数据集
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

4.4 模型训练与评估

训练模型这一步，说白了就是让模型自己学习数据里的规律。我个人习惯设置100个epoch，配合早停机制。

# 训练模型
history = model.fit(
    X_train, y_train,
    epochs=100,
    validation_split=0.2,
    callbacks=[keras.callbacks.EarlyStopping(patience=10)]
)

# 评估模型
test_loss, test_acc = model.evaluate(X_test, y_test)
print(f'测试集准确率: {test_acc:.4f}')

你想想看，为什么要加早停？因为训练时间长了，模型可能会过拟合。早停就是当验证集准确率不再提升时，自动停止训练。我在项目中试过，不加早停，模型准确率反而会下降。

4.5 模型量化与转换为.tflite

这一步最关键。STM32的资源有限，不能直接跑完整的TensorFlow模型。我们需要把模型量化，转成.tflite格式。

为什么要量化？说白了就是把模型里的浮点数变成整数。比如原来用float32，现在用int8。这样模型体积能缩小4倍，推理速度也能提升不少。

转换代码：

# 转换为TFLite模型
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# 启用量化
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 设置代表数据集进行校准
def representative_dataset():
    for i in range(100):
        yield [X_train[i:i+1].astype(np.float32)]

converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

tflite_model = converter.convert()

# 保存模型
with open('model_quantized.tflite', 'wb') as f:
    f.write(tflite_model)

嗯，这里要注意。代表数据集这一步不能省。它用来校准量化参数，直接影响量化后的模型准确率。我曾经偷懒没做校准，结果量化后的模型准确率从95%掉到了60%，简直惨不忍睹。

4.6 验证量化模型

转换完之后，一定要验证一下。看看量化后的模型还能不能正常工作。

# 加载量化模型
interpreter = tf.lite.Interpreter(model_path='model_quantized.tflite')
interpreter.allocate_tensors()

# 获取输入输出信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 推理测试
input_data = X_test[0].astype(np.int8)
interpreter.set_tensor(input_details[0]['index'], [input_data])
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])

print(f'预测结果: {np.argmax(output_data)}')
print(f'真实标签: {y_test[0]}')

4.7 小结

这一章我们走完了从模型设计到量化的完整流程。说白了就是三步：搭模型、训练、量化。每一步都有坑，但踩过之后你就记住了。

下一章，我们会把量化好的.tflite模型部署到STM32上。到时候你就知道，原来AI离嵌入式这么近。

好了，这一章就到这儿。有什么问题，欢迎交流。