CIFAR10 TFLM Quick Build Guide catalogs
I. Overview II. Environment Preparation 2.1 Virtual environment 2.2 Using Google Colab III. Core Steps 3.1 CIFAR10 dataset 3.2 Model creation 3.3 Model Training 3.4 Model Transformation 3.5 Inference Validation 3.6 Benchmark Performance 3.7 Full Implementation 3.7 Simplified version (Colab) IV. TFLite Micro Deployment V. Quick Verification Application Examples VII. Conclusion Reference
Spoiler (Highlight to read) more details, please see the attachment. more details, please see the attachment.
I. Overview
CIFAR-10 : Alex Krizhevsky, University of Toronto The CIFAR-10 public dataset, and one of the most classic and commonly used entry-level benchmark datasets in computer vision, contains 60,000 32x32 color images (50,000 training, 10,000 testing) in 10 categories, such as airplanes, cars, birds, cats, and more.
tflm_cifar10 : Demonstrates how to run a CIFAR-10 image classification model in real-time on a microcontroller from NXP using the TensorFlow Lite Micro framework. That is, a pre-trained convolutional neural network model for the CIFAR-10 dataset is deployed to the MCU, giving it the ability to recognize 10 classes of common objects (airplanes, cars, birds, cats, etc.).
Model: a lightweight CNN model with 3 convolutional layers, ReLU activation layer, pooling layer and a fully connected layer.
Input: 32x32 pixel color image.
Output: probability that an image belongs to one of the 10 categories in CIFAR-10.
This document: provides the complete process of building for the CIFAR10 dataset, from the dataset, model training transformation, deployment of inference of the rapid implementation of the program, can be used as an example of the tflm_cifar10 (inference-based) front complement, this paper does not involve the deployment and optimization of the end-side.
...
"""
CIFAR10 快速训练、测试、部署与推理完整流程
"""
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # 减少TensorFlow日志
import tensorflow as tf
import numpy as np
import time
import matplotlib.pyplot as plt
print(f"TensorFlow版本: {tf.__version__}")
print(f"NumPy版本: {np.__version__}")
class CIFAR10QuickPipeline:
def __init__(self):
"""初始化管道"""
self.model = None
self.tflite_model = None
def load_data(self, sample_size=1000):
"""加载简化数据集"""
print("\n1. 加载CIFAR10数据集...")
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
# 预处理
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
# 使用少量数据（快速训练）
x_train_small = x_train[:sample_size]
y_train_small = y_train[:sample_size]
x_test_small = x_test[:200]
y_test_small = y_test[:200]
# 转换为独热编码
y_train_onehot = tf.keras.utils.to_categorical(y_train_small, 10)
y_test_onehot = tf.keras.utils.to_categorical(y_test_small, 10)
print(f"训练数据: {x_train_small.shape}")
print(f"测试数据: {x_test_small.shape}")
return (x_train_small, y_train_onehot), (x_test_small, y_test_onehot)
def create_simple_model(self):
"""创建简化CNN模型"""
print("\n2. 创建简单CNN模型...")
model = tf.keras.Sequential([
# 输入层
tf.keras.layers.Input(shape=(32, 32, 3)),
# 卷积层1
tf.keras.layers.Conv2D(8, (3, 3), padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D((2, 2)),
# 卷积层2
tf.keras.layers.Conv2D(16, (3, 3), padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D((2, 2)),
# 全连接层
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dropout(0.2),
tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(
optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy']
)
model.summary()
self.model = model
return model
def train_model(self, x_train, y_train, x_test, y_test, epochs=10):
"""训练模型"""
print("\n3. 训练模型...")
# 回调函数：早停
callbacks = [
tf.keras.callbacks.EarlyStopping(
monitor='val_loss',
patience=3,
restore_best_weights=True
)
]
history = self.model.fit(
x_train, y_train,
epochs=epochs,
batch_size=32,
validation_data=(x_test, y_test),
callbacks=callbacks,
verbose=1
)
# 评估模型
test_loss, test_acc = self.model.evaluate(x_test, y_test, verbose=0)
print(f"\n测试准确率: {test_acc:.4f}")
return history
def convert_to_tflite(self):
"""转换为TFLite格式"""
print("\n4. 转换为TFLite格式...")
# 转换为TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
# 优化配置
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float32]
# 转换
tflite_model = converter.convert()
# 保存模型
with open('cifar10_model.tflite', 'wb') as f:
f.write(tflite_model)
# 保存为字节数组（用于嵌入式部署）
self.save_as_c_array(tflite_model)
self.tflite_model = tflite_model
model_size = len(tflite_model) / 1024
print(f"模型大小: {model_size:.1f} KB")
return tflite_model
def save_as_c_array(self, tflite_model):
"""保存为C数组格式"""
c_array = '// 自动生成的CIFAR10模型数组\n'
c_array += '#include \n\n'
c_array += 'const unsigned char cifar10_model_tflite[] = {\n'
# 每行显示12个字节
for i in range(0, len(tflite_model), 12):
line_bytes = tflite_model[i:i+12]
c_array += ' ' + ', '.join(f'0x{b:02x}' for b in line_bytes) + ',\n'
c_array += '};\n\n'
c_array += f'const unsigned int cifar10_model_tflite_len = {len(tflite_model)};\n'
with open('cifar10_model_array.h', 'w') as f:
f.write(c_array)
print("C数组已保存: cifar10_model_array.h")
def test_tflite_inference(self, x_test, y_test, num_tests=10):
"""测试TFLite推理"""
print(f"\n5. 测试TFLite推理 ({num_tests}个样本)...")
if self.tflite_model is None:
with open('cifar10_model.tflite', 'rb') as f:
self.tflite_model = f.read()
# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_content=self.tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# 类别名称
class_names = ['飞机', '汽车', '鸟', '猫', '鹿',
'狗', '青蛙', '马', '船', '卡车']
correct = 0
times = []
for i in range(min(num_tests, len(x_test))):
# 准备输入
input_data = x_test[i:i+1]
# 推理
start_time = time.perf_counter()
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
inference_time = time.perf_counter() - start_time
times.append(inference_time)
# 获取输出
output = interpreter.get_tensor(output_details[0]['index'])
predicted_class = np.argmax(output[0])
actual_class = np.argmax(y_test[i])
# 检查是否正确
if predicted_class == actual_class:
correct += 1
print(f"样本 {i+1}: 预测={class_names[predicted_class]:<5} "
f"实际={class_names[actual_class]:<5} "
f"时间={inference_time*1000:.1f}ms "
f"{'✓' if predicted_class == actual_class else '✗'}")
accuracy = correct / num_tests
avg_time = np.mean(times) * 1000
print(f"\n推理统计:")
print(f" 准确率: {accuracy:.1%} ({correct}/{num_tests})")
print(f" 平均推理时间: {avg_time:.1f}ms")
print(f" 推理速度: {1000/avg_time:.0f} FPS")
return accuracy, avg_time
def benchmark_performance(self, x_test):
"""性能基准测试"""
print("\n6. 性能基准测试...")
interpreter = tf.lite.Interpreter(model_content=self.tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
# 预热
test_input = x_test[0:1]
for _ in range(10):
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
# 基准测试
num_runs = 100
start_time = time.perf_counter()
for _ in range(num_runs):
interpreter.invoke()
total_time = time.perf_counter() - start_time
avg_time = total_time / num_runs * 1000
print(f"基准测试结果:")
print(f" 总推理次数: {num_runs}")
print(f" 总时间: {total_time*1000:.1f}ms")
print(f" 平均推理时间: {avg_time:.1f}ms")
print(f" 推理速度: {1000/avg_time:.0f} FPS")
return avg_time
def save_model_summary(self):
"""保存模型摘要"""
summary = []
self.model.summary(print_fn=lambda x: summary.append(x))
with open('model_summary.txt', 'w') as f:
f.write('\n'.join(summary))
f.write(f"\n\n模型信息:")
f.write(f"\n参数数量: {self.model.count_params():,}")
f.write(f"\n保存时间: {time.ctime()}")
print("模型摘要已保存: model_summary.txt")
def main():
"""主函数"""
print("=" * 60)
print("CIFAR10 快速训练、测试、部署管道")
print("=" * 60)
# 创建管道
pipeline = CIFAR10QuickPipeline()
# 1. 加载数据
(x_train, y_train), (x_test, y_test) = pipeline.load_data(sample_size=2000)
# 2. 创建模型
pipeline.create_simple_model()
# 3. 训练模型
history = pipeline.train_model(x_train, y_train, x_test, y_test, epochs=15)
# 4. 保存模型摘要
pipeline.save_model_summary()
# 5. 转换为TFLite
pipeline.convert_to_tflite()
# 6. 测试推理
pipeline.test_tflite_inference(x_test, y_test, num_tests=20)
# 7. 性能测试
pipeline.benchmark_performance(x_test)
print("\n" + "=" * 60)
print("流程完成！生成的文件:")
print(" 1. cifar10_model.tflite - TFLite模型")
print(" 2. cifar10_model_array.h - C数组格式")
print(" 3. model_summary.txt - 模型摘要")
print("=" * 60)
if __name__ == "__main__":
main()
...
VII. Conclusion
The purpose of this paper is to take a common image classification scenario (CIFAR10) as an example, so that readers can quickly understand the complete process from data building, model creation, training, inference and validation, which can be used as an example of the tflm_cifar10 (end-side inference-based) front complement, this paper does not involve the end-side deployment and optimization.
查看全文
