如果你用Pytorch訓練的模型那么可以參考我的步驟,使用的是Tensorflow的話參考官方文檔即可,但流程都是一樣的,每一步我都會提到部分操作細節及注意事項
官方教程
要詳細學習的話tflite-micro里的微控制器章節下都詳細看(頁面左側目錄最下方),要先嘗試跑通就直接跳到下一節內容跟著做,
一切的開始,環境踩坑提示(已經有模型不需要訓練可跳過)
如果是自己用搭建模型結構從0訓練,按正常模型構建、訓練流程即可,要用預訓練的模型的話可以參考官方步驟,本筆記僅為個人探索過程的記錄不詳細展開模型訓練過程了,最后要導出的模型為文件為.tflite格式的,建議不要用太高版本的tensorflow,詳情看tensorflow官方的tflite micor相關說明,關于版本可能會踩非常多坑,如果用pytorch的話會用到onnx進行轉化,也有版本的坑,所以友情提示把以下需要確認對應版本在官方文檔里找到,新建一個虛擬環境,然后先初步過一遍確定沒有版本沖突再開始自己的模型代碼開發,否則環境的坑會導致不停重來,一定要新建一個虛擬環境,不要偷懶用自己已有的環境!!!!
- GPU驅動
- CUDA版本
- pytorch/tensorflow(二選一,推薦tensorflow)和python版本
- Numpy版本
- onnx、onnxtf、onnxruntime版本(pytorch)
我用的是服務器,CUDA11.2,python3.9,主Pytorch ,本地的ESP-IDF版本是5.4.1,這里貼一個我的環境版本(太長了放在另一篇筆記里),
模型訓練注意事項
官方文檔里有這樣的提示
使用預訓練模型
如果要用預訓練模型的話一定要參考tflite-micro官方文檔確定可支持的模型有哪些,并不是在PC上能跑通就可以的,還要考慮在ESP32上的部署。
自己構建模型從0訓練
不管是Pytorch還是tensorflow框架,僅使用tflite-micro支持的算子,否則就要自己添加自定義算子,我碰到了這個坑感興趣可以看我的這篇筆記【踩坑隨筆】TensorFlowLite_ESP32庫中不包含REDUCE_PROD算子,手動移植
預處理注意事項
一定要確定模型的輸入輸出以及圖片的預處理,我的模型輸入為[1,3,64,64], 轉化部署模型考慮采用uint8量化,所以在訓練環節的預處理我直接不做0~1的歸一化,而是采用[0,255],以下是我的預處理代碼,重點是最后的歸一化操作用transforms.Lambda(lambda x: x * 255)
而不用Normalize
def get_transforms():train_transform = transforms.Compose([transforms.Resize((72, 72)),transforms.RandomCrop(64),transforms.RandomHorizontalFlip(0.5),transforms.RandomRotation(10),transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),transforms.ToTensor(),# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])transforms.Lambda(lambda x: x * 255) # 轉為 [0,255]])val_transform = transforms.Compose([transforms.Resize((64, 64)),transforms.ToTensor(),# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])transforms.Lambda(lambda x: x * 255)])return train_transform, val_transform
模型導出轉換
Tensorflow模型可以直接轉存tflite格式,前面提到我采用的是Pytorch導出的是pth格式,所以走pth–>onnx->tensorflow模型–>tflite模型的路線,這個思路是ai生成的,可能會有冗余,感覺饒了很大的彎子,等我探索完有更好的方式再來更新
pth轉tflite(INT8量化)
參數里的路徑只是示例,改成你自己的對應文件路徑即可,model_path
是訓練完保存好的pth
格式的模型,data_path
是數據集的文件夾,output_dir
是轉換完保存tflite格式的模型的文件夾,重點注意converter.inference_input_type = tf.uint8
,這決定了部署模型的輸入
def pytorch_to_tflite(model_path="./models/best_model.pth", data_path="./datasets49",output_dir="./models/tfmodels"):"""將PyTorch模型轉換為TensorFlow Lite模型"""print("Starting model conversion process...")# 1. 加載PyTorch模型print("Loading PyTorch model...")checkpoint = torch.load(model_path, map_location='cpu')num_classes = checkpoint['num_classes']model = ESPNetV2(num_classes=num_classes)model.load_state_dict(checkpoint['model_state_dict'])model.eval()print(f"Model loaded with {num_classes} classes")# 2. 準備校準數據print("Preparing calibration data...")_, val_transform = get_transforms()val_dataset = CustomImageDataset(data_path, transform=val_transform, train=False)val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)# 3. 導出為ONNXprint("Exporting to ONNX...")dummy_input = torch.randn(1, 3, 64, 64)onnx_path = os.path.join(output_dir, "model.onnx")torch.onnx.export(model,dummy_input,onnx_path,export_params=True,opset_version=11,do_constant_folding=True,input_names=['input'],output_names=['output'],dynamic_axes={'input': {0: 'batch_size'},'output': {0: 'batch_size'}})print(f"ONNX model saved to: {onnx_path}")# 4. 使用 onnx-tf 轉換為 TensorFlow SavedModelprint("Converting ONNX to TensorFlow using onnx-tf...")try:subprocess.check_call([sys.executable, "-m", "pip", "install", "onnx-tf"])from onnx_tf.backend import prepareexcept Exception as e:print("Failed to import or install onnx-tf:", e)returnonnx_model = onnx.load(onnx_path)tf_rep = prepare(onnx_model)tf_model_dir = output_dirtf_rep.export_graph(tf_model_dir)print(f"TensorFlow SavedModel saved to: {tf_model_dir}")# 5. 創建校準數據生成器def representative_dataset():for i, (data, _) in enumerate(val_loader):if i >= 100: # 只使用100個樣本進行校準breakyield [data.numpy().astype(np.float32)]# 6. 轉換為TensorFlow Lite (INT8量化)print("Converting to TensorFlow Lite with INT8 quantization...")converter = tf.lite.TFLiteConverter.from_saved_model(tf_model_dir)converter.optimizations = [tf.lite.Optimize.DEFAULT]converter.representative_dataset = representative_datasetconverter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] # 確保形狀操作兼容converter.inference_input_type = tf.uint8converter.inference_output_type = tf.uint8 tflite_model = converter.convert()# 保存TFLite模型tflite_path = os.path.join(output_dir, "model_quantized.tflite")with open(tflite_path, 'wb') as f:f.write(tflite_model)print(f"Quantized TFLite model saved to: {tflite_path}")print(f"Model size: {len(tflite_model) / 1024:.2f} KB")# # 評估轉換后的模型# print("\nEvaluating converted model...")# evaluate_tflite_model(tflite_path, data_path)return tflite_path
提取tflite模型的量化參數和算子
這一步只是為了方便做驗證,如果對模型結構和轉化過程非常確定的話可以不提取,具體代碼看【淺學】從tflite模型提取算子和量化參數,會提取到以下信息,主要用來在部署中做驗證,包括輸入輸出格式、量化參數和采用的算子
{"input": {"name": "serving_default_input:0","shape": [1,3,64,64],"dtype": "<class 'numpy.uint8'>","scale": 1.0,"zero_point": 0},"output": {"name": "PartitionedCall:0","shape": [1,51],"dtype": "<class 'numpy.uint8'>","scale": 0.060637399554252625,"zero_point": 150}
}
inline tflite::MicroMutableOpResolver<8> CreateModelResolver() {tflite::MicroMutableOpResolver<8> resolver;// 注冊 CONV_2Dmicro_op_resolver.AddConv2d();// 注冊 DEPTHWISE_CONV_2Dmicro_op_resolver.AddDepthwiseConv2d();// 注冊 FULLY_CONNECTEDmicro_op_resolver.AddFullyConnected();// 注冊 MEANmicro_op_resolver.AddMean();// 注冊 PADmicro_op_resolver.AddPad();// 注冊 QUANTIZEmicro_op_resolver.AddQuantize();// 注冊 RESHAPEmicro_op_resolver.AddReshape();// 注冊 TRANSPOSEmicro_op_resolver.AddTranspose();return resolver;
}
上面提取生成的這個算子文件不能直接用哈,寫法不對,但是算子是對應的,需要對應上tflite-micro支持的算子,只是大小寫不一樣的話實際是同一個算子部署的代碼用對的寫法就行,但是如果提取到的算子在tflite-micro里沒有就要自定義了,參考tflite官方文檔
tflite轉C數組
模型訓練好導出.tflite格式后,可以通過執行以下命令轉換成C數組,這個在Windows系統上直接用Git bash切換到model.tflite文件夾下,然后執行這個命令(或者安裝了git配有git環境直接在終端里也可以)要注意文件路徑不能有空格和中文,然后我們就能得到一個model.cc
文件,里面包含一個數組和一個表示數組大小的常量,確認一下數組大小,如果明顯比例程的大很多可能會導致Flash不夠用
$ xxd -i model.tflite > model.cc
圖片格式轉換
準備好你要進行圖像分類的10張圖片重命名為image0~image9,自己記錄好這10張圖片對應的類別方便你驗證最終推理結果對不對,然后進行轉換,轉換成二進制文件,一定要注意在轉換過程就做好了預處理,跟訓練環節的保持一致,convert_image_to_binary_pytorch(image_path, output_path, target_size=(64,64))
這個函數里的對應你自己的輸入和預處理進行修改
import os
import glob
from PIL import Image
import numpy as np
import torchvision.transforms as transformsdef convert_image_to_binary_pytorch(image_path, output_path, target_size=(64,64)):"""使用 PyTorch transform 處理圖片并保存為二進制文件確保和訓練/驗證輸入一致"""# 打開圖片并轉換為 RGBimg = Image.open(image_path).convert('RGB')# 定義 transform,和驗證集一致transform = transforms.Compose([transforms.Resize(target_size),transforms.ToTensor(), # float32, [0,1]transforms.Lambda(lambda x: x * 255) # float32, [0,255]])img_tensor = transform(img) # C,H,W, float32img_tensor = img_tensor.byte().numpy() # uint8# 保證是 NCHWimg_tensor = img_tensor.astype(np.uint8)# 保存二進制文件with open(output_path, 'wb') as f:f.write(img_tensor.tobytes())# 打印信息,便于檢查print(f"Converted {os.path.basename(image_path)} -> {os.path.basename(output_path)}")print(f"Shape: {img_tensor.shape}, dtype: {img_tensor.dtype}, min/max: {img_tensor.min()}/{img_tensor.max()}\n")return output_pathdef convert_folder_to_binary_pytorch(folder_path, output_folder=None):"""批量將文件夾下圖片轉換為二進制文件"""# 支持的圖片格式image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp']if output_folder is None:output_folder = folder_pathos.makedirs(output_folder, exist_ok=True)for ext in image_extensions:for image_path in glob.glob(os.path.join(folder_path, ext)):output_name = os.path.splitext(os.path.basename(image_path))[0]output_path = os.path.join(output_folder, output_name)try:convert_image_to_binary_pytorch(image_path, output_path)except Exception as e:print(f"Failed to convert {image_path}: {e}")# 使用示例
if __name__ == "__main__":folder_path = "orignal_img"output_folder = "images"if os.path.isdir(folder_path):convert_folder_to_binary_pytorch(folder_path, output_folder)else:print("錯誤:指定的路徑不是一個有效的文件夾")
模型和數據驗證
寫個腳本批量執行,分別使用轉換前的pth模型和tflite模型、原始圖片和轉換的二進制圖片進行組合推理驗證,如果原始模型原始圖片的推理結果有誤識別不用管,那也是正常的是模型準確率的問題,確保這四種組合的輸出結果都是一致的(誤識別的也是一樣的輸出),否則根據這四個推理結果進行對比去查是哪一個環節出的問題,同一原始圖片pth模型和tflite模型輸出一致確保轉換的tflite模型沒問題,tflite模型使用原始圖片和轉換的二進制圖片輸出一致說明轉換的圖片沒有問題
代碼部署
例程創建
步驟參照這篇vscode+ESP-IDF+ESP32S3N16R8跑通TensorFlow Lite Micro for Espressif Chipsets的hello_word例程,把例程hello_world
換成
person_detection
即可,先跑通這個例程確保你的環境和硬件都沒有問題
idf.py create-project-from-example "esp-tflite-micro:person_detection"
修改宏用圖片測試
例程直接運行成功使用攝像頭獲取圖片進行推理,然后我們選擇用本地二進制圖片測試,打開eap_main.h
,取消#define CLI_ONLY_INFERENCE 1
的注釋
// Enable this to do inference on embedded images
#define CLI_ONLY_INFERENCE 1
運行成功的話在串口的終端輸入detect_image 0 (0~9的任意數字都可以),就可以看到檢測結果了。!!!!!!以下的前提是先把例程跑通哈!例程的README,md里有例程操作說明
模型加載
然后把前面轉化的model.cc復制到main
文件夾下,打開,第一行添加頭文件#include "model.h"
,把model_tflite[]
數組類型改成下面這樣,只改類型其他不動
#include "model.h"alignas(8) const unsigned char model_tflite[] = {....};
const int model_tflite_len = 12345;
然后再main文件夾下新建一個model.h,輸入以下代碼
#ifndef MODEL_H
#define MODEL_Hextern const unsigned char model_tflite[];
extern const int model_tflite_len;#endif
然后打開main文件夾下的CMakeLists.txt,SRCS后面添加“model.cc”
idf_component_register(SRCS"detection_responder.cc""image_provider.cc""main.cc""main_functions.cc""model_settings.cc""person_detect_model_data.cc""app_camera_esp.c""esp_cli.c""model.cc"# PRIV_REQUIRES console static_images spi_flashPRIV_REQUIRES console test_images spi_flashINCLUDE_DIRS "")
打開main文件夾下的main_functions.cc,找到void setup()
函數,修改模型加載,把原來加載的例程的模型g_person_detect_model_data
換成了我們自己的模型model_tflite
void setup() {// Map the model into a usable data structure. This doesn't involve any// copying or parsing, it's a very lightweight operation.// model = tflite::GetModel(g_person_detect_model_data);model = tflite::GetModel(model_tflite);
然后往下滑找到算子注冊的部分,替換成你自己的算子,比如我改成我自己用到的算子(跟前面提取的是對應的)
模型設置
打開main文件夾的model_settings.cc
,把kCategoryLabels[kCategoryCount]
數組中的類別替換為你的類別
打開main文件夾的model_settings.h
,把下面的參數修改為你的輸入尺寸,比如我的是(1,3,64,64),對應NCHW,我的類別總共有51類,輸出為(1,51),對應修改以下內容
constexpr int kNumCols = 64; //W
constexpr int kNumRows = 64; //H
constexpr int kNumChannels = 3; //Cconstexpr int kCategoryCount = 51; //類別
圖片數據替換
找到static_images文件夾中的sample_iamges文件夾,把里面的image圖片文件刪掉,把我們前面轉換好的自己的沒有后綴的image0~image9復制到這個文件夾下
推理代碼修改
打開main文件夾下的main_functions.cc
文件,void run_inference(void *ptr)
函數里第一個#if
到最后一個#endif
的內容不動,其他內容刪掉,然后修改成以下代碼,結合注釋自行理解一下,跟PC端的推理操作其實是一樣的步驟
void run_inference(void *ptr) {memcpy(input->data.uint8, ptr, input->bytes);#if defined(COLLECT_CPU_STATS)long long start_time = esp_timer_get_time();
#endif// Run the model on this input and make sure it succeeds.if (kTfLiteOk != interpreter->Invoke()) {MicroPrintf("Invoke failed.");}#if defined(COLLECT_CPU_STATS)long long total_time = (esp_timer_get_time() - start_time);printf("Total time = %lld\n", total_time / 1000);//printf("Softmax time = %lld\n", softmax_total_time / 1000);printf("FC time = %lld\n", fc_total_time / 1000);printf("DC time = %lld\n", dc_total_time / 1000);printf("conv time = %lld\n", conv_total_time / 1000);printf("Pooling time = %lld\n", pooling_total_time / 1000);printf("add time = %lld\n", add_total_time / 1000);printf("mul time = %lld\n", mul_total_time / 1000);/* Reset times */total_time = 0;//softmax_total_time = 0;dc_total_time = 0;conv_total_time = 0;fc_total_time = 0;pooling_total_time = 0;add_total_time = 0;mul_total_time = 0;
#endifTfLiteTensor* output = interpreter->output(0);float output_probs[kCategoryCount];float sum = 0.0f;// 量化 uint8 -> float (反量化)for (int i = 0; i < kCategoryCount; i++) {output_probs[i] = (output->data.uint8[i] - output->params.zero_point) * output->params.scale;sum += expf(output_probs[i]); // softmax}// 計算 softmax 概率static float max_val = output_probs[0];for (int i = 1; i < kCategoryCount; i++) {if (output_probs[i] > max_val) max_val = output_probs[i];}float sum_exp = 0.0f;for (int i = 0; i < kCategoryCount; i++) {output_probs[i] = expf(output_probs[i] - max_val); // 防止 exp 溢出sum_exp += output_probs[i];}int max_idx = 0;float max_prob = 0.0f;for (int i = 0; i < kCategoryCount; i++) {output_probs[i] /= sum_exp;if (output_probs[i] > max_prob) {max_prob = output_probs[i];max_idx = i;}}int category_score_int = (max_prob) * 100 + 0.5;MicroPrintf("Detected: %s, score: %d%%",kCategoryLabels[max_idx], category_score_int);;
}
到這里就可以編譯運行了,前提是先把例程跑通哈!例程的README,md里有例程操作說明!!!!