[PyTorch] 모델 시각화 툴 세가지 - Torchviz, HiddenLayer, Netron (Model visualization)

PyTorch 모델을 시각화할 수 있는 툴 세가지를 소개한다. 

 

출처: https://stackoverflow.com/questions/52468956/how-do-i-visualize-a-net-in-pytorch

How do I visualize a net in Pytorch?

 

 

1. Torchviz

https://github.com/szagoruyko/pytorchviz

 

backward pass를 이용해 모델을 시각화한다. 중간에 어떻게 gradient를 계산하고 backward 연산을 진행하는지 모두 보여준다.

 

Install:

pip install torchviz

 

Usage:

from torchviz import make_dot
make_dot(y_pred, params=dict(model.named_parameters())

 

Example:

 

 

 

2. HiddenLayer

https://github.com/waleedka/hiddenlayer

 

Install:

pip install hiddenlayer

 

Usage:

import hiddenlayer as hl
hl.build_graph(model, input)

 

Example:

 

기본적인 사용법은 위와 같은데 그래프 색깔을 바꾸거나, 여러 node들을 합쳐 하나의 block으로 만든다거나 하는 customizing option들이 있다. 

그런데 연산그래프에 너무 불필요한 부분이 많긴 하다. BN을 constant-cast-add-div 의 과정으로 나타낸다던지..

Hiddenlayers는 모델 시각화 이외에도 training 과정에서 사용할 수 있는 다양한 시각화 기능들을 제공한다.

 

 

3. Netron

https://github.com/lutzroeder/netron

 

desktop application으로 다운로드 받을 수도 있고, browser version을 이용할 수도 있다.

PyTorch 모델도 experimental support를 진행 중이라고 해서 torch.save로 저장한 모델과 state_dict를 각각 넣어봤는데 둘다 좀.. 유의미한 정보가 아니다.

 

model.state_dict를 불러왔을 때

model을 통째로 불러왔을 때

 

ONNX format으로 모델을 저장한 후 불러오면 잘 보인다는 것 같다. 내 모델은 complex number을 input으로 받는데 ONNX format이 complex format을 지원을 안한다고 해서 그냥 안해봤음.

그러나 일단 browser로도 사용할 수가 있고, zoom이나 이동 등 UI가 잘 되어 있어 사용하기 매우 편한 것 같다.

 

---

 

내가 사용한 모델은 complex를 input으로 받는 모델인데, ONNX format이 complex를 지원하지 않아서 Netron은 .pth 파일로만 해볼 수 있었고, hiddenlayer 역시 중간에 ONNX로 변환하는 과정을 거쳐서 사용을 못해봤다.

결론은 가장 간단한 것은 Torchviz, customizing이 가능하고 기능이 많은 것은 HiddenLayer, browser-base라 편리한 것은 Netron인듯 ❗