yolo v5 导出 Android 手机上可以运行的 TFLite 模型

更新日期: 2023-03-29 阅读次数: 2937 字数: 560 分类: AI

TFLite 是 TensorFlow Lite 的缩写。TFLite 在 Android 上的运行速度更快,因为它专门针对移动设备进行了优化。但是,最后我发现我实在找不到在 Android 上执行 yolov5 TFLite 模型的示例,转而使用 pytorch mobile 的 torchscript lite 模型实现了在 Android 上的目标检测

虽然没有在 Android 上用 tflite 模型实现功能,但是记录一下导出 tflite 模型的操作方法。

导出命令

python export.py --weights yolov5s.pt --include tflite
  • 目前支持 tflite 格式的导出,即导出 TensorFlow Lite 格式的 model,后缀为 yolov5s.tflite
  • 如果是要导出 ios 需要使用的 model 格式,则是 CoreML。

例如:

python export.py --weights /content/drive/MyDrive/yolo/wheat_head_best.pt --include tflite

导出时,生成的 tflite 会很贴心的位于与原 model 同级目录。

TensorFlow SavedModel: export success ✅ 9.1s, saved as /content/drive/MyDrive/yolo/wheat_head_best_saved_model (27.0 MB)

TensorFlow Lite: starting export with tensorflow 2.11.0...
WARNING:absl:Found untraced functions such as tf_conv_2_layer_call_fn, tf_conv_2_layer_call_and_return_conditional_losses, tf_conv_3_layer_call_fn, tf_conv_3_layer_call_and_return_conditional_losses, tf_conv_4_layer_call_fn while saving (showing 5 of 268). These functions will not be directly callable after loading.
2023-03-28 05:10:04.153800: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:362] Ignored output_format.
2023-03-28 05:10:04.153854: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:365] Ignored drop_control_dependency.
TensorFlow Lite: export success ✅ 52.4s, saved as /content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite (13.5 MB)

Export complete (64.0s)
Results saved to /content/drive/MyDrive/yolo
Detect:          python detect.py --weights /content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite
Validate:        python val.py --weights /content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '/content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite')
Visualize:       https://netron.app

验证

找几张图片测试一下导出的 tflite 模型是否可用。

!python detect.py --weights /content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite --source /content/drive/MyDrive/yolo/wheat_heads

detect: weights=['/content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite'], source=/content/drive/MyDrive/yolo/wheat_heads, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5 🚀 v7.0-131-g0c8de3f Python-3.9.16 torch-1.13.1+cu116 CUDA:0 (Tesla T4, 15102MiB)

2023-03-28 05:22:04.899803: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-28 05:22:06.005299: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/lib/python3.9/dist-packages/cv2/../../lib64:/usr/lib64-nvidia
2023-03-28 05:22:06.005429: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/lib/python3.9/dist-packages/cv2/../../lib64:/usr/lib64-nvidia
2023-03-28 05:22:06.005451: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Loading /content/drive/MyDrive/yolo/wheat_head_best-fp16.tflite for TensorFlow Lite inference...
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
image 1/4 /content/drive/MyDrive/yolo/wheat_heads/微信图片_20230326171007.jpg: 640x640 103 persons, 652.2ms
image 2/4 /content/drive/MyDrive/yolo/wheat_heads/微信图片_20230326171020.jpg: 640x640 35 persons, 445.7ms
image 3/4 /content/drive/MyDrive/yolo/wheat_heads/微信图片_20230326171034.jpg: 640x640 152 persons, 462.9ms
image 4/4 /content/drive/MyDrive/yolo/wheat_heads/微信图片_20230326171053.jpg: 640x640 59 persons, 460.6ms
Speed: 1486.9ms pre-process, 505.3ms inference, 6.9ms NMS per image at shape (1, 3, 640, 640)
Results saved to runs/detect/exp

参考

  • 官方文档:https://github.com/ultralytics/yolov5/issues/251

微信关注我哦 👍

大象工具微信公众号

我是来自山东烟台的一名开发者,有感兴趣的话题,或者软件开发需求,欢迎加微信 zhongwei 聊聊, 查看更多联系方式

tags: pytorch yolov5