YoloV4 到 TFLite 模型给出完全错误的预测

Question

我在自定义数据集上训练了我的 yolov4 模型，它给了我大约 90-100% 的准确率，太棒了！但我想将权重转换为 TFlite，因为我需要在边缘设备上执行一些检测任务。但是在我将其转换为第一个 TF 权重和 TFLite 权重之后，它会输出真正随机的预测，如下图所示。第一张图片来自 Darknet 检测，第二张来自 TFLite 检测。

暗网检测使用命令：

!./darknet detector test data/obj.data cfg/yolov4-obj.cfg /content/drive/MyDrive/yolov4/backup/yolov4-obj_last.weights /content/drive/MyDrive/yolov4/test/t2.webp -thresh 0.3

使用 TFLite 权重：python detect.py

我不知道是什么原因造成的。此外，在将暗网权重转换为 tflite 权重时，我没有遇到任何错误。以下是脚本：

1. save_model.py - 保存自定义 yolov4 tf 模型以进行 tflite 转换

python save_model.py --weights ./data/custom.weights --output ./checkpoints/custom-416 --input_size 416 --model yolov4 --framework tflite

import tensorflow as tf
from absl import app, flags, logging
from absl.flags import FLAGS
from core.yolov4 import YOLO, decode, filter_boxes
import core.utils as utils
from core.config import cfg

flags.DEFINE_string('weights', './data/yolov4.weights', 'path to weights file')
flags.DEFINE_string('output', './checkpoints/yolov4-416', 'path to output')
flags.DEFINE_boolean('tiny', False, 'is yolo-tiny or not')
flags.DEFINE_integer('input_size', 416, 'define input size of export model')
flags.DEFINE_float('score_thres', 0.2, 'define score threshold')
flags.DEFINE_string('framework', 'tf', 'define what framework do you want to convert (tf, trt, tflite)')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')

def save_tf():
  STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config()

  input_layer = tf.keras.layers.Input([FLAGS.input_size, FLAGS.input_size, 3])
  feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
  bbox_tensors = []
  prob_tensors = []
  if FLAGS.tiny:
    for i, fm in enumerate(feature_maps):
      if i == 0:
        output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      else:
        output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      bbox_tensors.append(output_tensors[0])
      prob_tensors.append(output_tensors[1])
  else:
    for i, fm in enumerate(feature_maps):
      if i == 0:
        output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      elif i == 1:
        output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      else:
        output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      bbox_tensors.append(output_tensors[0])
      prob_tensors.append(output_tensors[1])
  pred_bbox = tf.concat(bbox_tensors, axis=1)
  pred_prob = tf.concat(prob_tensors, axis=1)
  if FLAGS.framework == 'tflite':
    pred = (pred_bbox, pred_prob)
  else:
    boxes, pred_conf = filter_boxes(pred_bbox, pred_prob, score_threshold=FLAGS.score_thres, input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
    pred = tf.concat([boxes, pred_conf], axis=-1)
  model = tf.keras.Model(input_layer, pred)
  utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
  model.summary()
  model.save(FLAGS.output)

def main(_argv):
  save_tf()

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass

1. python convert2tflite.py - 现在，将上述 yolov4 tf 模型转换为 tflite 权重。

import tensorflow as tf
from absl import logging
import numpy as np

def save_tflite(model_path, tf_weights_path, quantize_mode='float16'):
  converter = tf.lite.TFLiteConverter.from_saved_model(tf_weights_path)
  tflite_model = converter.convert()
  open(model_path, 'wb').write(tflite_model)

  logging.info("model saved to: {}".format(model_path))

def demo(model_path):
  interpreter = tf.lite.Interpreter(model_path=model_path)
  interpreter.allocate_tensors()
  logging.info('tflite model loaded')

  input_details = interpreter.get_input_details()
  print(input_details)
  output_details = interpreter.get_output_details()
  print(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[i]['index']) for i in range(len(output_details))]

  print(output_data)

def main():
  save_tflite(
    tf_weights_path='./checkpoints/anpr-416',
    model_path='./checkpoints/anpr-416.tflite'
  )
  demo(model_path='./checkpoints/anpr-416.tflite')

if __name__ == '__main__':
    try:
        main()
    except SystemExit:
        pass

如果有人能提供任何见解，我将不胜感激！谢谢！

Answer 1

好吧，我做错的是，在转换为位于core/config.py 内的 TF 模型时，我没有更改 obj.names 文件中的类名，因此在 obj.names 文件中进行了所需的更改之后在core/config.py 内部，我现在能够做出正确的检测！