tflm逃课！为什么我选择TinyMaix，一步步踩坑跑通第一个demo
ChatGLM的官方环境低于我的配置，我暂时也不想更改，所以大模型剪枝量化暂时放一下。我很早就关注Tinyml的发展，曾尝试搭建纯c语言的裸奔tinyml工程，没有计划性，现在想和工业界的工程接轨。于是尝试tensorflowlite_micro.
Tensorflowlite-micro是Google旗下的ai搭建框架Tensorflow专门为嵌入式边缘设备设计的一个工具，在工程界得到了广泛的安全认可。所以我最先尝试的也是这个。和别人不同，我没有使用bazel环境，所以基本都是调用官方工具手搓了整个部署。以下细讲。
首先是环境搭建，就遇到了大问题。
Tensroflow官方鼓励的搭配是Tensorlflow2.15版本，配合内部集成的keras库，但是要调用剪枝方法，必须要一个新的工具库tensorflow_model_optimization。然而，我安装0.8版tensorflow_model_optimization的时候，每次都出现keras无法导入的情况，多次修改"TF_USE_LEGACY_KERAS"没有很好的效果（官方考虑到兼容旧版，特意搞了这个参数，结构导致旧版keras导入和新版不一样）。
最后，我尝试的组合是官方最新版的Tensorflow 2.21.0+tensorflow_model_optimization0.8+tf_keras,这个tf_keras可以看作官方为了兼容旧版搞的一个兼容包，可能会在未来优化掉。亲测可以剪枝量化，这个demo就是这个组合。
首先是官方的第一个示例hello world,也就是一个正弦函数的神经网络。于是我用Tensorflow的工具自己打建了一个训练神经网络。
import numpy as np
import tensorflow as tf
import tensorflow_model_optimization as tfmot
import matplotlib.pyplot as plt
import os
# ====================== 环境配置 ======================
os.environ["TF_USE_LEGACY_KERAS"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
# ====================== 1. 准备数据 ======================
x_train = np.linspace(0, 2 * np.pi, 1000, dtype=np.float32)
y_train = np.sin(x_train)
x_test = np.linspace(0, 2 * np.pi, 200, dtype=np.float32)
y_test = np.sin(x_test)
# ====================== 2. 构建基础模型 ======================
def build_model():
inputs = tf.keras.Input(shape=(1,), name='input_layer')
x = tf.keras.layers.Dense(32, activation='relu', name='dense_1')(inputs)
x = tf.keras.layers.Dense(32, activation='relu', name='dense_2')(x)
outputs = tf.keras.layers.Dense(1, name='output_layer')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.compile(optimizer='adam', loss='mse')
return model这里分类很清晰了，就是随机生成一些噪声，然后用内置的Dense构建连接层，整个连接层的结构是：(FC表示连接层)
输入层（dim 1*1*1）->FC1(32)->FC2(32)->输出层(1*1*1)
之后先训练模型:
# ====================== 3. 先正常预训练 ======================
print("Stage 1/4: Pre-training base model...")
base_model = build_model()
base_model.fit(x_train, y_train, epochs=100, batch_size=32, verbose=1)
base_loss = base_model.evaluate(x_test, y_test, verbose=0)
print(f"Pre-training done. Test Loss: {base_loss:.6f}")
##输出Pre-training done. Test Loss: 0.002890接下来是剪枝:
# ====================== 4. 应用剪枝 ======================
print("\nStage 2/4: Applying pruning...")
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
batch_size = 32
epochs = 50
end_step = np.ceil(len(x_train) / batch_size).astype(np.int32) * epochs
pruning_params = {
'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
initial_sparsity=0.0,
final_sparsity=0.5, # 剪枝 50%，可调整
begin_step=0,
end_step=end_step
)
}
model_for_pruning = prune_low_magnitude(base_model, **pruning_params)
model_for_pruning.compile(optimizer='adam', loss='mse')微调：
# ====================== 5. 剪枝后微调 ======================
print("\nStage 3/4: Fine-tuning after pruning...")
callbacks = [tfmot.sparsity.keras.UpdatePruningStep()]
model_for_pruning.fit(
x_train, y_train,
epochs=epochs,
batch_size=batch_size,
callbacks=callbacks,
verbose=1
)
pruned_loss = model_for_pruning.evaluate(x_test, y_test, verbose=0)
print(f"Pruning + fine-tuning done. Test Loss: {pruned_loss:.6f}")
final_model = tfmot.sparsity.keras.strip_pruning(model_for_pruning)
##Pruning + fine-tuning done. Test Loss: 0.002294最后量化，一定要先剪枝后量化，因为如果先量化，意味着模型已经损失了大量精度，在此基础上剪枝，误差会非常大。(不过工业上有时会进行结构化剪枝再量化，这里对于大多数情况，特别是非结构剪枝)
可以看一下图像:(用matlab的)
可以看到，拟合效果还是不错的
测试输入x=pi/2,输出y=0.95;
接下来应用Tensorflow自带的工具，保存成.tflite文件。这是一个二进制文件，保存了模型参数等信息。
with open('sin_micro_model_pruned.tflite', 'wb') as f:
f.write(tflite_quant_model)这样我们就得到了sin_micro_model_pruned.tflite文件大小是5kb，之后所有的模型参数转换由此展开。
但是,Tensorflowlite-micro为了工业界的灵活性和高度兼容性，提供的是interepter接口管理模型。一方面加大了模型加载的开销，另一方面，导致底层更多的函数，算子（c的c++的汇编的）放在不同的地方，单看文件结构非常混乱，不适合初学者。所以我退而求其次，选择了国内社区的TinyMaix。这是一个国人开发的专注于mcu的边缘ai算子库项目，精简了tensorflowlite-micro的大部分，保留主要框架。https://github.com/sipeed/TinyMaix.非常适合初学者快速了解内部原理。
我们知道，深度学习简单来说就是很多神经网络层，所以只要我们想Tensorflow或者其他神经网络架构那样，把层封装起来，按照参数调用，也可以实现相关功能（这就是为什么我写个c库裸奔神经网络也可以，但是没有系统的算子库，导致没法移植）。而这个项目就是实现简单的算子库和一些常用芯片的硬件加速比如arm的各种指令架构。
项目的核心文件就5个，一个是基本类型的定义文件，和模型函数加载的头文件：tinymaix.h
这个文件主要包含的就是模型各种参数的定义，比如layer 层，会记录是不是最后一层，如果是输出结构，不是就继续运行神经网络。比如input层，需要输入维度，比如我的是sin函数那就是1维.
所有的类型定义都在这里。包括模型加载api:
* @brief 加载模型
* @param mdl 模型句柄
* @param bin 模型二进制数据指针（Flash中的const数组）
* @param buf 静态工作缓冲区（用户预分配的SRAM数组）
* @param cb 层运行回调函数（不需要传NULL）
* @param in 返回输入张量信息
* @return 错误码，TM_OK为成功
*/
tm_err_t tm_load (tm_mdl_t* mdl, const uint8_t* bin, uint8_t*buf, tm_cb_t cb, tm_mat_t* in);
/**
* @brief 卸载模型
* @param mdl 模型句柄
*/
void tm_unload(tm_mdl_t* mdl);
/**
* @brief 输入数据预处理
* @param mdl 模型句柄
* @param pp_type 预处理类型
* @param in 输入原始数据
* @param out 预处理后的数据
* @return 错误码
*/
tm_err_t tm_preprocess(tm_mdl_t* mdl, tm_pp_t pp_type, tm_mat_t* in, tm_mat_t* out);
/**
* @brief 运行模型推理
* @param mdl 模型句柄
* @param in 输入张量
* @param out 输出张量
* @return 错误码
*/
tm_err_t tm_run (tm_mdl_t* mdl, tm_mat_t* in, tm_mat_t* out);
(原作者这里没有太多注释，我自己加的)几个比较重要的：首先，需要定义一个tm_mat_t input结构体,存放输入的维度，输入的值。还要定义一个输出的结构体tm_mat_t output,存放输出的值。模型的句柄也要先定义出来tm_mdl_t mymdl，通过tm_load填充，类似于类的创建对象。这里的：
* @param bin 模型二进制数据指针（Flash中的const数组）就是转化后的数组的指针，至于转化，我等一下还要在后面讲一下几个坑。
tm_cb_t是回调函数，可以放一些异常处理的钩子，如果不在意检查，可以填NULL.注意，uint8_t*buf如果填NULL,会自动调用malloc函数，对于我的单片机stm32f103c8t6,频繁malloc很容易造成碎片内存导致崩溃，所以对于资源紧张状况考虑用静态内存。也就是先分配buf的缓冲区，然后把指针给函数。这样，模型的加载就完成了。
好了，第二个是和移植相关的关键文件：tm_port.h
里面定义了关于架构平台的宏，要根据自己的需要开启，还有fpu是否启用，我的是f1 系列没有fpu,所以关闭，配置如下:
/******************************* 架构平台定义 ************************************/
#define TM_ARCH_CPU (0) //default, pure cpu compute
#define TM_ARCH_ARM_SIMD (1) //ARM Cortex M4/M7, etc.
#define TM_ARCH_ARM_NEON (2) //ARM Cortex A7, etc.
#define TM_ARCH_ARM_MVEI (3) //ARMv8.1: M55, etc.
#define TM_ARCH_RV32P (4) //T-head E907, etc.
#define TM_ARCH_RV64V (5) //T-head C906,C910, etc.
#define TM_ARCH_CSKYV2 (6) //cskyv2 with dsp core
#define TM_ARCH_X86_SSE2 (7) //x86 sse2
#define TM_OPT0 (0) //default, least code and buf
#define TM_OPT1 (1) //opt for speed, need more code and buf
#define TM_OPT2 (2) //TODO
/******************************* 核心端口配置（适配STM32F1+int8模型） ************************************/
// 平台架构：STM32F1为Cortex-M3，无DSP/FPU，使用纯CPU计算
#define TM_ARCH TM_ARCH_CPU
// 优化等级：优先最小代码和内存占用，适配STM32F1有限资源
#define TM_OPT_LEVEL TM_OPT0
// 模型类型：固定使用int8量化模型，和训练转换环节匹配
#define TM_MDL_TYPE TM_MDL_INT8
// 快速缩放：无FPU芯片开启可提升速度，int8模型精度损失可忽略
#define TM_FASTSCALE (1)
// 本地数学函数：使用内置快速近似算法，避免调用标准库libm，解决指数函数计算错误问题
#define TM_LOCAL_MATH (1)
// 模型统计功能：关闭以节省Flash和SRAM，量产场景无需开启
#define TM_ENABLE_STAT (0)
// 动态内存分配：彻底关闭，使用用户预分配的静态内存，避免malloc跑飞
#define TM_ENABLE_MALLOC (0)
// 文件系统支持：彻底关闭，避免fopen调用，模型直接从Flash数组加载
#define TM_ENABLE_FILE (0)
// 模型资源上限：适配你的正弦小模型，大幅缩减预留内存
#define TM_MAX_CSIZE (16) // 最大通道数，你的模型仅16个隐藏层神经元
#define TM_MAX_KSIZE (1*1) // 最大卷积核尺寸，你的全连接模型无需卷积
#define TM_MAX_KCSIZE (16) // 最大卷积核*通道数，适配你的小模型( 以上有一些注释是自己加的，原来没有，有出入以原项目为准)
基本看看注释应该理解在说什么。
然后是主要是关于性能和时间的检测函数，如果不追求性能，只是跑通，可以不管。但是为了移植和以后，还是讲一下。
/******************************* DBG TIME CONFIG ************************************/
#include <sys/time.h>
#include <time.h>
#define TM_GET_US() ((uint32_t)((uint64_t)clock()*1000000/CLOCKS_PER_SEC))
#define TM_DBGT_INIT() uint32_t _start,_finish;float _time;_start=TM_GET_US();
#define TM_DBGT_START() _start=TM_GET_US();
#define TM_DBGT(x) {_finish=TM_GET_US();\
_time = (float)(_finish-_start)/1000.0;\
TM_PRINTF("===%s use %.3f ms\n", (x), _time);\
_start=TM_GET_US();}<sys/time.h>应该是为了兼容linxu等平台。主要是clock()函数的实现，裸机是没有这个的。所以我们可以改写一下clock(),比如摘到系统的Systick Handle钩子，进行一个时间运算，当成clock函数。基本思路:(先注释掉sys/time.h)
#ifndef __TIME_H
#define __TIM_H
#include "stm32f10x.h"
uint32_t clock(void);
#endif
#include "time.h"
#include "stm32f10x.h"
volatile uint32_t systick_ms = 0;
void SysTick_Handler(void) {
systick_ms++;
}
void SysTick_Init(void) {
SystemCoreClockUpdate();
SysTick_Config(SystemCoreClock / 1000); // 1ms 中断
}
uint32_t clock(void){
return systick_ms;
}
这样就实现了简易的裸机时钟。
ok，基本的配置差不多就在这里。然后是关于硬件加速的文件：
可以看到支持arm_mvei,neon,simd等指令加速。指令加速的实现基本是采用官方的加速指令集，比如:
#if TM_MDL_TYPE==TM_MDL_INT8
#define PARALLEL_CNT 16
TM_INLINE void tm_dot_prod(mtype_t* sptr, mtype_t* kptr,uint32_t size, sumtype_t* result)
{
uint32_t cnt;
int8x16_t vecA;
int8x16_t vecB;
sumtype_t sum = 0;
cnt = size/PARALLEL_CNT;
while (cnt > 0) {
vecA = vld1q_s8(sptr);
vecB = vld1q_s8(kptr);
sum = vmladavaq(sum, vecA, vecB);
cnt--; sptr += PARALLEL_CNT; kptr += PARALLEL_CNT;
}
cnt = size%PARALLEL_CNT;
if (cnt > 0U) {
mve_pred16_t p0 = vctp8q(cnt);
vecA = vld1q_s8(sptr);
vecB = vld1q_s8(kptr);
sum = vmladavaq_p(sum, vecA, vecB, p0);
}
*result = sum;
return;
}这个vld1q_s8就是可以让cpu在一个时钟周期进行多次乘加运算，对于矩阵乘法位构建基础的神经网络来说，非常有效。不然原本就是一个时钟周期可能搬数据到寄存器，一个时钟周期运算，效率低了好几倍。
然后是模型层函数实现和模型加载的实现文件:tm_stat.c和tm_layer.c。
到了这一步，移植的框架就搭好了。
接下来就是找到我们的“血肉”——模型参数。我们刚从Tensorflow到导出的。tflite文件是可以兼容的，因为TinyMaix本身脱胎于Tensorflow的量化，所以在资源极度紧张的情况下，可以直接无缝衔接Tensorflow的量化剪枝。
(实测的时候，Tensroflow有时会出现一些路径检测错误的问题，所以转化模型权重我放到了linux环境里。TinyMaix官方要求的配置为numpy和tensorflow 2.14.0,实际运行文件还要一个Pillow,用于图像。虽然这个项目不带图像，但是由于图像的逻辑放在了同一个py文件下，所以最好安装一下，不然可能报错。)
找到官方的tools文件夹,找到tflite2tmdl，官方给出了运行参数:
Usage: python3 tflite2tmdl.py tflite_name tmdl_name mdl_type out_deq in_dims out_dims我用的时候：
python -m tools.tflite2tmdl sin_micro_model_pruned.tflite model.tmdl int8 1 1,1,1 1
tflite_name:输入的tflite文件路径
tmdl_name:输出tmdl,.h的文件路径
mdl_type:模型类型，我的是量化后的int8.
out_deq:输出反量化使能，我的是1
in_dims:输入维数我的是1
out_dims:输出维数，我的是1
之后会看到输出了model.h,model.tmdl在工作目录下，把它们转移到自己的工作目录下。
好了，这下文件都有了。.h文件应该长这样:
#ifndef __MODEL_FILE__H
#define __MODEL_FILE__H
#include <stdint.h>
#define MDL_BUF_LEN (64)
#define LBUF_LEN (1408)
const uint8_t mdl_data[1872]={\包括了缓冲区长度，数组等信息。虽然默认是const,还是提一下。const会把这个数组当成静态变量，从而在flash分配内存。如果没有就会在sram分配。对于模型权重文件这种比较大的，建议在flash分配。可以通过配置生成.map 文件查看内存分配情况。如果参数位置在0x08000000(对于stm32f103 的flash默认)，说明是正确的，在 flash分配。如果分配在0x2000000(对于stm32f103默认)说明分配到了sram上，看需求改。可以看到量化剪枝之后模型的文件参数下降了很多。
好的，这下万事具备。开始编写主函数：
#include "tm_port.h"
#include "tinymaix.h"
#include "model.h"
#include "LED.h"
static uint8_t tm_buf[MDL_BUF_LEN] __attribute__((aligned(4)));
static tm_mdl_t mymdl;输入之后用preprocess,因为模型文件自己保存了量化参数，不需要手写。输出因为开启了反量化，所以不用自己再写一遍。
这里aligned(4)是要保证4字节对齐，不然后面内存会混乱。后面也会有调试讲一下这个东西。
接下来的代码就比较简单了：
int main(void)
{
LED_Init();
tm_err_t ret;
// 1. 加载模型，获取输入张量描述
tm_mat_t input;
ret = tm_load(&mymdl, mdl_data, tm_buf, NULL, &input);
if (ret != TM_OK) while(1);
// 3. 准备浮点输入张量
float f_input_val = 1.507;
tm_mat_t input_fp = input;
input_fp.dataf = &f_input_val;
// 4. 准备量化输入缓冲区（大小根据输入形状确定）
int input_size = input.h * input.w * input.c;
int8_t q_input_buf[input_size];
tm_mat_t input_q = input;
input_q.data = q_input_buf;
// 5. 执行预处理：浮点 -> 量化（自动使用模型量化参数）
ret = tm_preprocess(&mymdl, TMPP_FP2INT, &input_fp, &input_q);
if (ret != TM_OK) while(1);
// 6. 运行推理
tm_mat_t output;
ret = tm_run(&mymdl, &input_q, &output);
if (ret != TM_OK) while(1);
float f_output_val = output.dataf[0];
while(1) {
Set_LEDIntensity(f_output_val);
}
}tm_load 函数传入你的输入结构体，输出结构体的地址给tm_run处理。
这里为了保险，把输入的形状写出来。LED是为了用PWM稍微指示一下对错，实际也可以用keil 的debug功能看参数。这里我用keil。
api调用部分，到这里就结束了。但是接下来才是踩坑的地方。我会展示一下keil调试遇到的问题，争取帮大家梳理一下数据流动的过程。 主要函数tm_run
进入keil debug。在tm_run打上断点:
如果连main都进不去，可以看一下call back窗口：
如果看到sys fileopen等字眼，可以看一下魔术棒里面microlib库开了没有，也可以看一下有没有使用静态内存，有可能是malloc导致的跑飞。
正常运行到tm_rum,可以关注一下的参数：输入输出。可以在watch窗口打开:
原本是空的，需要自己手动键入f_input_val( 输入等)。检测这些值。正常情况，输入是1.507左右的float(pi/2)。输出是0(初始化为0)。
好的，进入函数:
重点关注_out.data,中间的数据流都在这里
可以在keil里面通过memory观察:
这是没初始化的时候。memory窗口位置和call back一样
tm_err_t TM_WEAK tm_run(tm_mdl_t* mdl, tm_mat_t* in, tm_mat_t* out)
{
tm_mat_t _in, _in1, _out;
tm_err_t res = TM_OK;
int out_idx = 0;
memcpy((void*)&_in, (void*)in, sizeof(tm_mat_t));
mdl->layer_body = mdl->b->layers_body;
for(mdl->layer_i = 0; mdl->layer_i < mdl->b->layer_cnt; mdl->layer_i++){
tml_head_t* h = (tml_head_t*)(mdl->layer_body);
if(mdl->layer_i>0) {
_in.data = (mtype_t *)(mdl->buf + h->in_oft);
memcpy((void*)&_in, (void*)(h->in_dims), sizeof(uint16_t)*4);
}
_out.data = (mtype_t *)(mdl->buf + h->out_oft);
memcpy((void*)&_out, (void*)(h->out_dims), sizeof(uint16_t)*4);函数的逻辑并不太复杂，首先要把你的输入in给到函数内部，第一次layer是0,不进入if,跳到下面:
_out.data = (mtype_t *)(mdl->buf + h->out_oft);此时，_out.data数据段的第一层应该输出的位置会填充成0（第1个数到第32个数）.这里mdl->buf就是你给的静态缓冲区首地址，h->oft就是该层输出相对的偏移值。对于第一层，这个结果是0，所以执行完后:
一行30个可以看到，前32个位置都是0，接下来运行完FC这里就会保存第一层的输出：
看看范围，输出正常。接下来要把这一层的输出变成下一层的输入：
代码段就是一个交接过程，直接看内存变化，第一层的输出整体后移了32位，为前面空出来，放第二层的输出：
可以看到，原来第一层输出的80 80 80 80（位置0到4）跑到了第二行（位置32到65）。
马上，前面的位置要被第二行的输出挤占，同样由：
_out.data = (mtype_t *)(mdl->buf + h->out_oft);完成。
这样，我们的神经网络就运行完了第二层FC，最后是FC->输出（1 维）。
此时我们查看一下h->in_oft
是0,那么可以预测，输出的应该是第一个位置?
也就是0x91(先别急，接下来就是内存的问题了)可以看到输出的数不止一个，第一行都占满了，但是我们只要一个，所以其他的都是脏数据。不过答案不是0x91.
嗯，还要回过头看一下模型怎么控制什么时候输出？用的就是layer结构体的is_out参数，如果当前层已经是最后，比如模型总共4层，当前第4层，那么is_out就会是有效，此时进入输出的分支:
if(res != TM_OK) return res;
if(mdl->cb) ((tm_cb_t)mdl->cb)(mdl, h); //layer callback
if(h->is_out) {
memcpy((void*)(&out[out_idx]), (void*)(&(h->out_dims)), sizeof(uint16_t)*4);
if(mdl->b->out_deq == 0 || TM_MDL_TYPE == TM_MDL_FP32) //fp32 do not need deq
out[out_idx].data = (mtype_t*)(TML_GET_OUTPUT(mdl, h));
else {
int out_size = h->out_dims[1]*h->out_dims[2]*h->out_dims[3];
float* outf = (float*)(TM_ALIGN(TML_GET_OUTPUT(mdl, h) + out_size));
for(int i=0; i<out_size; i++) //do dequant
outf[i] = TML_DEQUANT(h, (TML_GET_OUTPUT(mdl, h))[i]);
out[out_idx].dataf = outf;
}
out_idx += 1;
}
mdl->layer_body += (h->size);最上面两行的是错误判断的钩子。
好的，然后又是memcpy,把输出维度的部分先放到out里面，对于我们是01 01 01 01,但是在内存里，由于out的结构体定义:
typedef struct{
uint16_t dims;
uint16_t h;
uint16_t w;
uint16_t c;
union {
mtype_t* data;
float* dataf;
};
}tm_mat_t;可以看到，是8位对齐的，所以我们在内存监控里看一下：
看，前八位:01 00 01 00 01 00 01 00就是照顾到内存对齐。
后八位：98 00 00 20这个涉及到stm32架构的小端问题，对应的是数据的地址，也就是0x20000098看看这个地址是啥？
由于我们直接反量化了，直接看out结构体的union里面 float地址的解析out.dataf。更直接的是之间看out[0].dataf[0]的值
。至此，整个数据的流动我们大概就有感觉了。
总结一下，Tensorflowlite-micro确实是业界的主选，因为它完整的生态，灵活性，可验证性，安全性无可替代。不过对于初学者想接触tinyml相关的工作，建议先从TinyMaix这样的项目开始，了解主干。然后去学习Tensorflow 这样的“重型武器”，了解它内存管理和算子设计的精妙。我相信它的interpretor设计和c++兼容接口有必要性。还有一点就是Tensorflow为了内存管理，增加了一些内存开销。对于TinyMaix可以用静态数组，代码总量也小。对于资源极度紧张的硬件也是不错的选择。
附注：TinyMaix在直接keil编译可能会报很多警告，主要是参数定义未使用，可以void来避免大量警告。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。
更多推荐
10
0- 0
已为社区贡献2条内容
所有评论(0)