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)

 

끝!

 

 

 

 

출처 :  https://bekusib.tistory.com/210