How to convert pt to tflite (yolov5 to tensorflow lite) tflite로 변환
업데이트:
이걸로 해결
github.com/zldrobit/yolov5/blob/tf-android/models/tf.py
$ python models/tf.py --weights weights/best.pt --cfg models/yolov5s.yaml --img-size 320
yolov5로 학습된 모델을 1차적으로 Android에서 사용하고자 한다.
그럴러면 최종적으로 .pt 파일을 .tflite로 변환해야한다.
그게 한번에 되냐?
아쉽게도 pt -> onnx -> pb -> tflite 순으로 변환해야한다.
명령어 하나로 되었다면 얼마나 좋았을까..
그래서 그 과정에서 꽤 많은 삽질을 하게 되었다.
잠깐 언급하자면 tensorflow 버전과 onnx 버전의 호환성..
버전업되면서 바뀐것도 많고 알 수 없는 에러 또한 많다.
이 분야에선 너무 흔한 특징이라,
이 글에선 삽질 과정은 간단하게만 언급하겠다.
- Win10에서 python 특정 모듈의 버전을 찾을 수 없다고 해서 conda에 python3.6을 설치하려니 설치가 안되었다,
- 그래서 MacOS(카탈레나)에서 yolov5를 받아 export하려니 'google::protobuf' 관련 이유로 에러가 발생.
- 그래서 Ubuntu에서 yolov5를 받아 모듈을 설치하려니 torch 설치 실패..
그래서 최종적으로 google colab에서 작업하기로 결정!
tensorflow 2.4.0과
pip를 이용해 설치되는 모듈의 최신버전 onnx 1.6.0 onnx-tf 1.7.0으로 작업을 시작하니
pt -> onnx는 정상적으로 잘 변환 되었다.
그런데 onnx -> pb부분이 onnx-tf 버전에 따라 나오는 아웃풋 파일 양식이 달랐다.
onnx-tf 1.7.0의 경우 디렉토리가 생성되며 그 안에 saved_mode.pb가 만들어지며,
onnx-tf 1.6.0에선 지정된이름.pb 파일이 생성된다.
그에 따라 문제가 발생되거나 pb -> tflite를 위한 명령이나 코드가 달라진다.
아무튼 1.6.0이나 1.7.0이나 둘다 변환이 되질 않았다.
이걸로 정말 온갖 방법을 다 찾아보면서 시도했지만, 잘 해결되지 않았다.
onnx -> pb 부분에서 파일이 잘 만들어지긴 했는데 속이 이상해서 pb -> tflite가 안되는걸까? 등
변환 과정이 많아 의심되는 부분이 한두개가 아니었다.
결론은,
tensorflow (2.4.0 -> 2.3.0) -> 1.15.0
onnx (1.7.0 <-> 1.6.0) -> 1.8.0
onnx-tf (1.7.0 <-> 1.6.0) -> 1.8.0
torch 1.7.0
을 세팅으로 변환에 성공했다.
onnx-tensorflow : github.com/onnx/onnx-tensorflow/tree/tf-1.x
그러고보니 onnx>=1.8.0이 떡하니 적혀있다. 왜 첨엔 안됐을까.. tensorflow 2.4였어서 그런가?? (이젠 다시 확인 귀찮..)
yolov5/requirements.txt
최종 파일 구조
최종 변환된 best.tflite 파일 확인 (netron.app/)
model properities
model architecture
참 과정과 다르게 결과 코드는 간단하다...
...
먼저 yolov5로 학습된 best.pt 파일을 준비한다.
yolov5에서는 기본적으로 onnx로 변환하는 코드를 제공한다. (개꿀~)
# pt to onnx
$ python /content/yolov5/export.py --weights best.pt --img 640 --batch 1
onnx-tf를 이용해 pb파일을 만든다.
# onnx to pb
import onnx
from onnx_tf.backend import prepare
onnx_model = onnx.load("/content/best.onnx") # load onnx model
tf_rep = prepare(onnx_model) # prepare tf representation
tf_rep.export_graph("/content/output/best.pb") # export the model
tf.lite를 이용해 .tflite 파일로 변환한다.
# pb to tflite
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
graph_def_file = "/content/output/best.pb"
input_arrays = ["images"]
output_arrays = ["output"]
tflite_file = "/content/output/best.tflite"
converter = tf.lite.TFLiteConverter.from_frozen_graph(graph_def_file, input_arrays, output_arrays)
tflite_model = converter.convert()
open(tflite_file, "wb").write(tflite_model)
끝!