我在自定义数据集上训练了我的yolov 4模型,它给了我大约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权重时没有遇到任何错误。下面是脚本:
save_model.py
-保存自定义yolov 4 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
python convert2tflite.py
-现在,将上面的yolov 4 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
如果有人能给予任何见解,我真的会很感激!谢谢!
2条答案
按热度按时间pengsaosao1#
好吧,我做错的是,我没有改变类的名称在
obj.names
文件,而转换到TF模型,这是位于core/config.py
,所以在obj.names
文件中所需的变化后,在core/config.py
,我现在能够作出正确的检测!hlswsv352#
我在yolo v5和tflite上也遇到了同样的问题(yolov5使用PyTourch权重,而不是暗网),你能解释一下你是如何解决这个问题的吗?