[PyTorch] torchsummary로 모델 구조 확인하기

torchsummary 패키지를 이용하면 pytorch로 작성한 모델의 구조를 간단하게 확인할 수 있다.

설치

pip install torchsummary

사용 방법

import torchsummary

# 모델 불러오기
model = UNet(1, 4)
# 요약하기 : (1, 240, 240)은 예시 입력 데이터 차원임
torchsummary.summary(model, (1, 240, 240))

출력 결과

잘린 윗 부분에는 모델 레이어 구조가 출력되며, 사진에 첨부된 것처럼 파라미터 수, 전체 모델 크기 등을 확인 가능하다.