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 聊聊， 查看更多联系方式
