[딥러닝] 딥러닝 모델을 간단하게 시각화하는 프로그램, Netron | Pytorch, ONNX

오늘은 자신이 직접 디자인한 딥러닝 모델을 시각화하는 작업을 해보겠습니다. 

 

우선, Pytorch로 구현된 모델을 ONNX로 저장해야 합니다. 

 

Pytorch로 ONNX로 저장하기 위해서 아래 명령어를 통해 설치합니다.

pip install onnx-pytorch

설치를 완료 후, 자신이 구현한 모델의 객체가 있다면, 아래 코드를 통해 모델을 저장합니다. 

import torch.onnx


input_names = ['Time series data']
output_names = ['Output']

x = torch.zeros(1, 10, 6).to(device)
torch.onnx.export(model, x, 'regression_mode.onnx', input_names=input_names, output_names=output_names)

 

이제 모델 시각화 프로그램인 Netron을 git을 통해 설치합니다. 

https://github.com/lutzroeder/netron

GitHub - lutzroeder/netron: Visualizer for neural network, deep learning, and machine learning models

설치를 하게 되면, 방금 저장한 모델이 아래와 같이 보입니다.

위 ONNX파일을 들어가게 되면, 시각화된 모델 아키텍처를 확인할 수 있습니다. export를 통해 쉽게 이미지 파일로 불러올 수 있고, 기본은 Vertical이라 세로로 길지만, Horizontal을 통해 가로로 시각화 할 수 있습니다.

Vertical

Horizontal