개발 관련/SW, App 관련

Darknet의 weights를 onnx로 컨버팅

소서리스25 2024. 8. 21. 11:15
반응형

지난번부터 계속 괴롭혀 왔던 onnx 컨버팅을 진행하게 되었다.

 

시간이 오래 걸린 이유는 고집이 좀 있었는데 darknet의 커스텀 데이터셋 훈련결과의 weights(가중치)를 바로 onnx로 컨버팅 하기 위한 방법을 찾기 위한 것이었기 때문이다.

 

또한 이를 unity 2018 버전에서 사용하기 위해서는 Yolo v3 이하에서만 사용할 수 있다. Unity 2018의 barracuda가 거기까지만 지원하기 때문이다. 거기다 python과 필요한 패키지의 지원 버전이 맞지 않으면 오류를 쏟아내고 그 원인을 해결하기 위한 방법 또한 구글링을 해도 각기 다 다른 원인이 많아 해결하기가 너무나 까다로웠다. 버전 특성과 환경상 거의 안된다고 보면 된다. 현재는 Yolo v8인가 까지 나온 상태라 더욱 그렇다.

 

그러니 방법을 찾는데 오래 걸리 수 밖에 없다.

 

아무튼 컨버팅 한 방법은 일단 weights 파일을 tensorflow의 keras를 통해 HDF5(.h5)로 컨버팅 한 다음 onnx로 다시 컨버팅 하는 것이다.

 

그러면 처음부터 정리하면 다음과 같다.

 

 

1. 가상환경 만들기

  - 깔끔한 상태로 시작해야 오류가 적어진다. 만약 여러 버전의 라이브러리를 설치하면 알수없는 충돌로 오류 발생 확률이 크다.

  - 명령어 : conda create -n tf python=3.7

 

 

2. 가상환경에 필요한 라이브러리 설치하기

1
2
3
4
5
6
7
  - pip install tensorflow==2.3.0
 
  - pip install onnx
 
  - pip install keras
 
  - pip install keras2onnx
cs

 

 

3. 컨버팅 도구 설치하기

 - 다음의 깃에서 도구 설치하기

GitHub - allanzelener/YAD2K: YAD2K: Yet Another Darknet 2 Keras

 

GitHub - allanzelener/YAD2K: YAD2K: Yet Another Darknet 2 Keras

YAD2K: Yet Another Darknet 2 Keras. Contribute to allanzelener/YAD2K development by creating an account on GitHub.

github.com

 

 

4. 가충치 파일 준비하기

  - .weights 파일

  - .cfg 파일

  - 해당 가중치 파일을 깃에서 다운받은 폴더에 복사해 두자

 

 

5. 일부 소스 수정하기

  - 참고한 사이트에서 64비트 환경에 맞춰줘야 한다고 한다.

1
2
3
4
> yad2k.py : weights_file.read를 16에서 20으로 변경
 
weights_header = np.ndarray(
        shape=(4, ), dtype='int32', buffer=weights_file.read(20))
cs

 

  - 참고한 사이트에 없지만 먼저 실행한 결과 'space_to_depth' 오류가 발생하는데 tensorflow 버전에 따라 다음을 참고하여 변경해 줘야 한다. yad2k의 기본이 tensorflow 버전 1.x인 듯..

1
2
3
4
5
6
* TensorFlow 1.x: tensorflow.space_to_depth
* TensorFlow 2.x: tensorflow.nn.space_to_depth
 
  > yad2k/models 폴더의 keras_yolo.py에서 space_to_depth 앞에 nn 붙여야 함.
 
  > tf.space_to_depth -> tf.nn.space_to_depth
cs

 

 

6. .h5로의 컨버팅 진행

 - python yad2k.py yolov2.cfg yolov2.weights yolov2.h5

 - 오류 없이 정상적으로 실행되면 같은 폴더에 yolov2.h5 파일이 생성됨

 

 

7. onnx 컨버팅 준비하기

  - 다음의 코드를 convert_keras_to_onnx.py로 저장한다.

1
2
3
4
5
6
7
8
9
10
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import load_model
import onnx
import keras2onnx
 
onnx_model_name = 'yolov2.onnx'
 
model = load_model('yolov2.h5')
onnx_model = keras2onnx.convert_keras(model, model.name)
onnx.save_model(onnx_model, onnx_model_name)
cs

 

 

8. python으로 컨버팅을 실행

  - python convert_keras_to_onnx.py

  - 그러면 최종적으로 onnx 파일이 생성된다. 이것을 Unity에서 로딩하면 된다.

 

 

9. Netron에서 확인해보기

- Netron

 

Netron

 

netron.app

 

- 확인결과 컨버팅은 잘 되었다.

Netron에서 확인
Netron에서 확인

 

 

* 참조사이트

YOLOv2 학습 결과(가중치)를 onnx 형식으로 변환 #ONNX - Qiita

 

YOLOv2の学習結果(weights)をonnx形式に変換する - Qiita

前の記事でtiny-YOLOv2の独自データセットでの訓練を行いましたが、このモデルをBarracudaで使用するためにはこの形式(weights/cfg)では対応できないため、ONNX形式に変…

qiita.com

 

 

반응형