深度学习模型部署框架之tflite

m0_62603533

2382人浏览 · 2024-12-08 18:25:56

m0_62603533 · 2024-12-08 18:25:56 发布

文章目录

0 前言
1 简介
2 模型准备
3 模型优化
4 利用tfl模型推理的demo
5 参考文档
6 edge-impulse
总结

0 前言

本文是深度学习模型部署框架的第二篇，更多部署框架可以查看本专栏的其他文章。

1 简介

概念：TensorFlow Lite 是TensorFlow在移动和嵌入式设备上的轻量级推理框架，专门针对资源受限的环境进行了优化，以便在手机、嵌入式设备上运行深度学习模型。
优势：缩减了模型的大小和功耗，适合嵌入式等资源受限的设备。
坑点：tfl 目前主要支持 cnn 相关的算子，对 rnn 等其他网络中的算子还没有很好的支持。对自家的tf框架支持比较好，对于其他框架支持不佳。

2 模型准备

2.1 导出模型为tfl格式

python api（推荐）：tf.lite.TFLiteConverter.from_saved_model()

import tensorflow as tf

# 转换模型
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # SavedModel目录的路径
tflite_model = converter.convert()

# 这里可以添加其他操作，比如量化、自定义tf操作支持等

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

命令行工具：tflite_convert

$ tflite_convert --help

`--output_file`. Type: string. Full path of the output file.
`--saved_model_dir`. Type: string. Full path to the SavedModel directory.
`--keras_model_file`. Type: string. Full path to the Keras H5 model file.

python库，可能不适用于某些特定类型的ONNX模型，参考onnx2tflite

from onnx2tflite import onnx_converter

res = onnx_converter(
        onnx_model_path = "/content/denoiser_model_fixed.onnx",
        need_simplify = True,
        output_path = "./models/",
        target_formats = ['tflite'],
    )

onnx->tf ，参考onnx_tf

# onnx_tf
import onnx
from onnx_tf.backend import prepare

onnx_model_path = "/content/denoiser_model_fixed.onnx"
onnx_model = onnx.load(onnx_model_path)

tf_rep = prepare(onnx_model)
tf_rep.export_graph("/content/denoiser_model.pb")

注意，如果要将其他格式的模型转为tfl，需要other format->onnx->tf->tfl，主要难点在于onnx->tf

2.2 导出时注意事项

动态形状的支持，tf转tflite时，由于tf.TensorListReserve操作需要在转换过程中具有静态的element_shape，但原生模型是动态的，所以在TensorFlow Lite 转换器中启用选择性TF操作支持，通过设置converter.target_spec.supported_ops = [ tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS ]实现。这里只是提到一个我实践中遇到的案例，更多的操作可以自定义。
环境配置：转换模型时环境的配置很容易出错，目前实践成功的环境配置为tensorflow=2.10.0，tensorflow-probability=0.16.0，protobuf=3.20.3，太低或太高会报错。
模型的opset版本：转onnx时，opset_version的设置会影响后续的过程，最终设置opset_version=12可以正常运行
当模型比较小的时候，导出为tfl时需要仔细验证模型前后的输出误差，一般来说模型越小，转换以为tfl后输出的误差越大，如果对误差不能接受，可以选择更保守一些的量化方案，如float16量化。

2.3 tfl模型检查

检验tfl的输出与原始模型的输出是否一致，或输出的误差是否可以容忍

3 模型优化

tfl提供训练后量化和量化感知训练两种量化形式，量化感知训练需要在模型训练阶段操作，所以下面主要讲解训练后量化，参考官网demo

tfl提供下面三种训练后量化方案：
在这里插入图片描述
方案选择指南：

3.1 训练后float16量化

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_quant_model = converter.convert()

3.2 训练后动态范围量化

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
# 启动默认的 optimizations 来量化所有固定参数(权重)
converter.optimizations = [tf.lite.Optimize.DEFAULT] 
tflite_quant_model = converter.convert()

3.3 训练后int8量化

import tensorflow as tf

# 定义生成代表性数据的函数
def representative_data_gen():
  for data in tf.data.Dataset.from_tensor_slices((audio)).batch(1).take(100):
    yield [tf.dtypes.cast(data, tf.float32)]
    
# 模型整型量化，输入输出还是浮点
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 启动默认的 optimizations 来量化所有固定参数(权重)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 要量化可变数据(模型输入/输出 和层之间的中间体)，提供 RepresentativeDataset，来估算所有可变数据的动态范围
converter.representative_dataset = representative_data_gen
tflite_model_quant = converter.convert()

# 全整型量化，自研vad采用的是这种量化方式
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]    # 先启动默认的optimizations将模型权重进行量化
converter.representative_dataset = representative_data_gen  # 使用代表数据集量化模型中间值
# 如果有任何的 ops不能量化，converter 将抛出错误
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] 
# 将输入和输出tensors设置为int8 类型
converter.inference_input_type = tf.int8   # or tf.uint8
converter.inference_output_type = tf.int8  # or tf.uint8

tflite_model_quant = converter.convert()

4 利用tfl模型推理的demo

import numpy as np
import tensorflow as tf  # tf==2.17.0
import librosa
import os

class VADInference:
    def __init__(self, model_path):
        # 模型参数
        self.SAMPLE_RATE = 16000
        self.DURATION = 1.0
        self.N_MELS = 40
        self.HOP_LENGTH = 512
        self.N_FFT = 2048

        # 加载TFLite模型
        self.interpreter = tf.lite.Interpreter(model_path=model_path)
        self.interpreter.allocate_tensors()

        # 获取输入和输出张量的细节
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()

    def preprocess_audio(self, audio_path, start=None, end=None):
        """预处理音频文件"""
        # 加载音频
        audio, _ = librosa.load(audio_path, sr=self.SAMPLE_RATE)

        # 确保音频长度一致
        if len(audio) < self.SAMPLE_RATE * self.DURATION:
            audio = np.pad(audio, (0, int(self.SAMPLE_RATE * self.DURATION) - len(audio)))
        else:
            audio = audio[int(16000*start):int(16000*end)]


        # 提取梅尔频谱特征
        mel_spec = librosa.feature.melspectrogram(
            y=audio,
            sr=self.SAMPLE_RATE,
            n_mels=self.N_MELS,
            n_fft=self.N_FFT,
            hop_length=self.HOP_LENGTH
        )

        # 转换为分贝单位
        mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)

        # 归一化
        mel_spec_db = (mel_spec_db - mel_spec_db.mean()) / mel_spec_db.std()

        # 添加batch和channel维度
        mel_spec_db = np.expand_dims(mel_spec_db, axis=[0, -1])

        # 确保数据类型匹配模型输入
        mel_spec_db = mel_spec_db.astype(np.float32)

        return mel_spec_db, audio

    def predict(self, audio_path, threshold=0.5, start=None, end=None):
        """对音频文件进行VAD预测"""
        # 预处理音频
        input_data, audio = self.preprocess_audio(audio_path, start=start, end=end)
        
        # 使用int8量化时，将float32输入转换为int8类型
        # input_scale, input_zero_point = self.input_details[0]["quantization"]
        # input_data = input_data / input_scale + input_zero_point  # input_data 除以量化的 input_scale，然后加上量化的 input_zero_point，使得输入数据恢复为浮动范围（即从 int8 类型恢复到原始的浮动数据）。
        # input_data = input_data.astype(self.input_details[0]["dtype"])  # 直接转换为int8

        # 设置输入张量，这里resize_tensor_input是为了适应不同长度的音频
        self.interpreter.resize_tensor_input(self.input_details[0]['index'], input_data.shape)
        self.interpreter.allocate_tensors()
        self.interpreter.set_tensor(self.input_details[0]['index'], input_data)

        # 运行推理
        self.interpreter.invoke()

        # 获取输出结果
        output_data = self.interpreter.get_tensor(self.output_details[0]['index'])
        
        # 将输出的int8转换为float32
        # output_details = self.interpreter.get_output_details()
        # print(output_details)
        # output = self.interpreter.get_tensor(output_details[0]['index'])
        # scale = output_details[0]['quantization_parameters']['scales']
        # zero_point = output_details[0]['quantization_parameters']['zero_points']
        # output_data = (output.astype(np.float32) - zero_point) * scale
        
        # 获取预测概率
        probability = output_data[0][0]

        # 根据阈值判断是否有人声
        has_voice = probability > threshold

        return {
            'probability': float(probability),
            'has_voice': bool(has_voice)
        }, audio

# 使用示例
if __name__ == "__main__":
    # 初始化VAD推理器
    vad = VADInference("/content/tflite_vad_model.tflite")


    audio_file = "/content/audio.mp3"
    result, audio = vad.predict(audio_path=audio_file, threshold=0.5, start=0, end=1)  # start=0, end=1是为了只取音频的前1s

    print(result)

5 参考文档

文档：https://www.tensorflow.org/lite/guide?hl=zh-cn
tfl model：https://www.tensorflow.org/lite/guide/hosted_models?hl=zh-cn
tfl tutorials：https://www.tensorflow.org/lite/tutorials?hl=zh-cn
参考示例的demo：https://www.tensorflow.org/lite/guide/android?hl=zh-cn、https://github.com/tensorflow/examples/tree/master/lite/examples
嵌入式部署：https://github.com/tensorflow/tflite-micro