如何将Yolov4-tiny模型权重转换为TFLite格式？

1. 准备必要文件 #

确保你已备好以下文件：

• 训练生成的权重文件：yolov4-tiny_custom_training_last.weights
• 对应修改后的模型配置文件：yolov4-tiny-custom.cfg
  • 将原配置中的width=416改为width=1024，height=416改为height=64
  • 修改classes值为你的目标类别数量（例如classes=3）
  • 找到每个[yolo]层上方的[convolutional]层，将filters值设为(类别数 + 5)*3（比如类别数为3时，filters=24）
• 类别名称文件：obj.names（每行填写一个类别名称）

2. 转换为TensorFlow SavedModel #

使用支持非正方形输入的转换脚本完成Darknet到TensorFlow格式的转换：

python save_model.py \
  --weights ./yolov4-tiny_custom_training_last.weights \
  --output ./saved_model \
  --input_size 64 1024 \
  --model yolov4-tiny \
  --tiny

注意：input_size参数顺序为(height, width)，需与cfg文件中的height=64、width=1024完全对应；若脚本要求顺序为(width, height)，则替换为1024 64。

3. 将SavedModel转为TFLite #

使用TensorFlow原生工具完成转换，确保支持非正方形输入：
创建convert_to_tflite.py脚本，内容如下：

import tensorflow as tf

# 加载SavedModel格式模型
converter = tf.lite.TFLiteConverter.from_saved_model('./saved_model')

# 启用默认优化（可选，按需开启）
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 支持Yolov4-tiny所需的TensorFlow自定义操作
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]

# 执行转换并保存TFLite模型
tflite_model = converter.convert()
with open('yolov4-tiny-custom.tflite', 'wb') as f:
    f.write(tflite_model)

运行脚本：

python convert_to_tflite.py

4. 验证TFLite模型可用性 #

通过TFLite Interpreter验证模型是否正常工作：

import tensorflow as tf
import numpy as np

# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path='yolov4-tiny-custom.tflite')
interpreter.allocate_tensors()

# 获取输入输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 生成测试输入（与模型输入尺寸一致的随机张量）
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)

# 执行推理
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

# 获取输出并打印形状（验证是否符合预期）
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data.shape)

若输出形状符合模型设计预期，则转换成功。

常见问题处理 #

• 输入形状不匹配报错：检查cfg文件的width/height与转换时指定的input_size是否完全统一，确保所有环节的输入尺寸一致。
• TFLite推理报错：确认转换时已启用SELECT_TF_OPS，Yolov4-tiny的部分层依赖TensorFlow自定义操作支持。

内容的提问来源于stack exchange，提问作者mj01