[ONNX] pytorch 모델을 tf, keras모델로 바꾸기

https://github.com/PINTO0309/onnx2tf

GitHub - PINTO0309/onnx2tf: Self-Created Tools to convert ONNX files (NCHW) to TensorFlow/TFLite/Keras format (NHWC). The purpos

 

pytorch에는 장점이 있다.

huggingface에서의 전이 모델들은 대부분 pytorch로 작성되어 있고, 모델의 개수도 많다.

 

그에 반해 keras에도 장점이 있다. 모델을 쌓거나 추가하기가 매우 쉽고, tpu를 호환하는 라이브러리가 많다는 것.

 

그렇기에 이를 호환하려면 pytorch를 keras로 바꾸거나, keras를 pytorch 모델로 바꿔야 한다.

이를 이어주는 중간다리가 onnx이다.

 

사실 torch에서 onnx로 바꾸는 건 쉽다. 자체적인 라이브러리를 제공한다.

하지만, onnx를 keras나 tf로 바꾸는 것 자체가 쉽지 않다.

사실상 2024-04-25 현재 필자가 모든 라이브러리를 시도해보았으나 유일하게 위에 있는 라이브러리만 가능했다.

 

심지어 colab 환경에서만 거의 가능하다. 다른 환경에서는 오류가 된다.

 

그러니 웬만하면 이를 시도하는 걸 추천한다.

 

아래는 필자가 직접 시도했던 코드이다. huggingface에서 모델을 불러오고, 그 모델을 convert하는 코드이다.

import torch
import onnxruntime
import numpy as np
import onnx2tf
import onnx
import tensorflow as tf
from tensorflow.lite.python import interpreter as tflite_interpreter
from transformers import Swinv2Config, Swinv2Model, AutoImageProcessor, AutoModelForImageClassification

model = Swinv2Model.from_pretrained("microsoft/swinv2-large-patch4-window12to16-192to256-22kto1k-ft")
dummy_input = torch.randn(1, 3, 224, 224)  # 입력 데이터 생성
torch.onnx.export(model, dummy_input, "swin_transformer.onnx", verbose=True)

 # Now, let's convert the ONNX model to TF
onnx2tf.convert(
    input_onnx_file_path="swin_transformer.onnx",
    output_folder_path="swin_transformer",
)

 

주의사항 : tflite은 pb나 다른 것으로 바꿀 시에 이미 정보 소실이 된 상태이므로, 원래 모델보다 성능이 떨어질 수 있다.

import tensorflow as tf

# TFLite 파일 로드
interpreter = tf.lite.Interpreter(model_path="your_model.tflite")
interpreter.allocate_tensors()

# TensorFlow 모델로 변환
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 입력 및 출력 텐서의 이름과 모양 추출
input_tensor_name = input_details[0]['name']
output_tensor_name = output_details[0]['name']
input_shape = input_details[0]['shape']
output_shape = output_details[0]['shape']

# TensorFlow 모델 생성
input_tensor = tf.keras.layers.Input(shape=input_shape[1:])
output_tensor = tf.keras.layers.Lambda(lambda x: interpreter.tensor(output_tensor_name)())(input_tensor)
tf_model = tf.keras.Model(inputs=input_tensor, outputs=output_tensor)

# 모델 요약 확인
tf_model.summary()

 

성능 평가는 해 보아야 겠으나, 사실상 선택지가 없다.

 

차라리 직접 만드는 게 더 나을 거 같은 기분이 들기도 한다.