4. 模型训练基础：TensorFlow/Keras入门、PyTorch入门、模型量化概念

好，咱们进入第四章。这一章，说白了就是给模型训练打地基。

很多做嵌入式的兄弟，一听到「训练」两个字就头大。觉得那是算法工程师的事。我刚开始也这么想，直到有一次在项目里被模型精度坑惨了——训练时好好的，部署到STM32上直接崩了。从那以后我明白，搞部署的人，必须懂一点训练。

这一章，我不打算讲太深的理论。咱们就聚焦三件事：TensorFlow/Keras怎么上手、PyTorch怎么上手、以及模型量化到底是个啥。这些都是你后面做STM32CubeAI部署时，绕不开的基础。

4.1 TensorFlow/Keras：快速原型的好帮手

我个人习惯，做快速验证的时候，首选Keras。为什么？因为它太友好了。你想想看，几行代码就能搭出一个神经网络，这对我们搞嵌入式的人来说，简直是福音。

咱们来看一个最简单的例子。假设你要做一个手写数字识别，用Keras怎么写？

import tensorflow as tf
from tensorflow import keras

# 1. 加载数据
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# 2. 数据预处理（归一化）
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# 3. 搭建模型
model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10, activation='softmax')
])

# 4. 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 5. 训练
model.fit(x_train, y_train, epochs=5, batch_size=32)

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

看到了吗？就这么几行。我在项目中遇到过很多次，用Keras搭原型，一天之内就能跑通。然后你再根据结果，决定要不要换更复杂的网络。

4.2 PyTorch：灵活且强大，适合深度定制

PyTorch呢？它比Keras更「底层」一点。说白了，就是给你更多的控制权。你想想看，Keras像自动挡，PyTorch像手动挡。搞嵌入式部署，有时候就需要手动挡的灵活性。

我记得第一次用PyTorch做模型量化时，发现它对动态图的支持特别好。调试起来非常直观。咱们还是用MNIST举例，看看PyTorch怎么写：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 1. 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 2. 定义网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleNet()

# 3. 定义损失和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练
for epoch in range(5):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1} 完成')

# 5. 测试
correct = 0
total = 0
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        _, predicted = torch.max(output.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

print(f'测试准确率: {100 * correct / total:.2f}%')

嗯，这里要注意。PyTorch的训练循环是显式写的，不像Keras那样一行model.fit()搞定。但好处是，你可以在循环里插入各种自定义操作，比如打印中间层的输出，或者做量化感知训练。

4.3 模型量化概念：把大模型塞进小芯片

好，终于到了重点——模型量化。你想想看，我们训练出来的模型，参数都是32位浮点数（FP32）。一个几MB的模型，在PC上跑没问题。但STM32这类MCU，Flash和RAM都只有几百KB到几MB，根本装不下。

怎么办？量化。说白了，就是把FP32的权重，用更少的位数来表示。最常见的是INT8量化，也就是8位整数。

数据类型	位数	内存占用（相对）	精度损失
FP32	32位	1x	无
FP16	16位	0.5x	极小
INT8	8位	0.25x	可接受（通常<2%）

从FP32到INT8，内存直接降到四分之一。一个4MB的模型，量化后只有1MB。这对STM32来说，简直是救命。

量化有两种常见方式：

• 训练后量化（Post-Training Quantization, PTQ）：模型训练完后，直接对权重做量化。简单快速，但精度可能下降。
• 量化感知训练（Quantization-Aware Training, QAT）：在训练过程中模拟量化效果，让模型学会适应低精度。精度更高，但训练时间更长。

我个人建议，如果你只是做快速部署，先用PTQ试试。如果精度掉得厉害，再上QAT。我在一个工业视觉项目中，PTQ后精度掉了1.5%，但用了QAT后只掉了0.3%。嗯，这个代价是值得的。

4.4 小结：选哪个框架？

好，咱们总结一下。TensorFlow/Keras和PyTorch，到底选哪个？

• 如果你追求快速验证，不想写太多代码，选Keras。它跟STM32CubeAI的兼容性也最好。
• 如果你需要深度定制，比如做量化感知训练，或者调试网络中间层，选PyTorch。它的灵活性更高。

至于模型量化，记住一句话：量化是嵌入式AI的必修课。没有量化，你的模型再牛，也跑不到STM32上。

下一章，咱们就正式进入STM32CubeAI的实战环节。到时候，你会看到这些训练好的模型，是怎么一步步变成能在芯片上跑的代码的。