[PyTorch] onnx을 이용한 모델 시각화 및 프레임워크 전환

ONNX(Open Neural Network Exchange)는 서로 다른 딥 러닝 프레임워크 간에 모델을 교환할 수 있도록 Microsoft와 Facebook에서 개발한 오픈 소스 프레임워크이다.

ONNX는 프레임워크를 바꿀 때 모델을 처음부터 다시 빌드하지 않고도 프레임워크 간에 모델을 전송할 수 있다.

또 프로그램을 사용해 모델을 시각화할 수 있다.

ONNX 사용의 주요 이점은 다양한 딥러닝 프레임워크의 장점을 활용할 수 있다는 것이다.

예를 들어 PyTorch에서 모델을 트레이닝하고 ONNX로 내보낸 다음 모델을 다시 트레이닝하지 않고도 TensorFlow와 같은 다른 프레임워크에서 같은 모델을 사용할 수 있다.

 

이제 ONNX을 사용해서 모델의 프레임워크를 전환하는 방법과 시각화하는 방법을 알아보자.

---

ONNX 설치

ONNX는 아래 명령어로 간단히 설치할 수 있다.

pip install onnx

간혹 protobuf 버전이 안 맞다고 에러를 뱉기는 하는데 사용하는데 문제가 없으므로 무시하자.

---

ONNX로 모델 내보내기

예시를 통해 살펴보자. PyTorch에서 모델을 내보내는 예시이다.

import torch
import torch.nn as nn
import torch.onnx

# Define your PyTorch model
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

model = MyModel()

# Export the PyTorch model to ONNX format
input_shape = (1, 1)
x = torch.randn(input_shape)
torch.onnx.export(model, x, "model.onnx", input_names=["input"], output_names=["output"])

맨 마지막 줄에 유의하자.

torch.onnx.export에 정의한 모델과 모델의 인풋을 넣으면 .onnx 파일을 만들 수 있다.

---

ONNX로 모델 가져오기

위에서 PyTorch 프레임워크에서 모델을 내보냈으니 TensorFlow에서 ONNX 모델을 가져오는 방법을 살펴보자.

import onnx
from onnx_tf.backend import prepare

# Load the ONNX model
onnx_model = onnx.load("model.onnx")

# Convert the ONNX model to TensorFlow format
tf_model = prepare(onnx_model)

# Use the TensorFlow model
output = tf_model.run({"input": x.numpy()})

이러면 모델을 다시 작성하지 않고도 PyTorch에서 만든 모델을 Tensorflow으로 가져올 수 있다.

---

ONNX 파일로 모델 시각화하기

Netron(https://netron.app/)과 같은 ONNX 형식의 파일을 지원하는 프로그램을 사용하여 ONNX 파일을 열거나 모델을 시각화할 수 있다.

 

Netron에 접속하자.

NETRON 웹

가운데 로고를 눌러 GitHub 페이지로 이동해 로컬에 NETRON을 설치할 수도 있고 Open Model ...을 눌러서 웹에서 작업할 수도 있다.

 

모델을 업로드하면 자동으로 시각화를 해준다.

.onnx 파일 업로드

왼쪽 위의 메뉴를 눌러 모양을 변경할 수 있고 png나 svg파일로 시각화한 자료를 다운로드 할 수 있다.

시각화한 파일을 다운로드

혹은 프로그램을 로컬에 설치했다면 다음 코드로 모델을 시각화할 수 있다.

import netron

# Load the ONNX model
model_path = "model.onnx"

# Start Netron web server and open the model in the default browser
netron.start(model_path)

그러면 브라우저에서 모델을 볼 수 있을 것이다.

---

살펴본 것처럼 ONNX는 프레임워크 간 모델을 주고받거나 모델을 시각화하는데 유용하게 쓸 수 있다.

 

사실 프레임워크를 전환할 일은 별로 없기 때문에 나는 ONNX를 발표 자료같은데 넣을 모델 시각화에 주로 사용한다.

모델이 간단하면 모델을 직접 그려도 되지만 모델이 복잡하고 크기가 클 경우 ONNX을 유용하게 사용할 수 있다.