[PyTorch] Model Structure & Parameters

모델의 구조 및 가중치를 확인하는 방법을 정리한다. 

 

torchsummary는 모델의 구조와 모델 변수를 간략히 알려주는 라이브러리이다. 

pip install torchsummary

별도의 설치가 필요하다.

 

import torch
from torch import nn
import torch.nn.functional as F
from torchsummary import summary

# GPU 또는 CPU 장치 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class Regressor(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(13, 50) # bias=True는 default
        self.fc2 = nn.Linear(50, 30)
        self.fc3 = nn.Linear(30, 1)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(F.relu(self.fc2(x)))
        x = F.relu(self.fc3(x))
        return x

모델을 보면 입력층에서 첫 번째 은닉층으로 넘어가는 부분의 가중치는 13 x 50 = 650개, 그리고 편향이 50개이므로 총 700개의 변수가 존재한다. 같은 방법으로 첫 번째 은닉층과 두 번째 은닉층 사이의 변수는 50 x 30 + 30 = 1530개, 두 번째 은닉층과 세 번째 은닉층 사이의 변수는 30 x 1 + 1 = 31개이다. 

 

# 모델 생성 및 디바이스로 이동
model = Regressor().to(device)

print(model)  # 모델 구조 출력

Regressor(
  (fc1): Linear(in_features=13, out_features=50, bias=True)
  (fc2): Linear(in_features=50, out_features=30, bias=True)
  (fc3): Linear(in_features=30, out_features=1, bias=True)
  (dropout): Dropout(p=0.5, inplace=False)
)

모델을 출력하면 __init__에서 정의된 구조를 확인할 수 있다. 

 

# 모델의 변수를 순서대로 출력
for parameter in model.parameters():
    print(parameter.size())

torch.Size([50, 13])
torch.Size([50])
torch.Size([30, 50])
torch.Size([30])
torch.Size([1, 30])
torch.Size([1])

model.parameters()를 통해 정의된 순서대로 변수를 얻을 수 있다. 

 

# 변수명에 직접 접근하여 출력
print(model.fc1.weight.size(), model.fc1.bias.size())

torch.Size([50, 13]) torch.Size([50])

변수명을 알고 있다면 변수명에 직접 접근하여 변수를 불러올 수 있다. 

 

# 변수명과 변수를 동시에 출력
for name, param in model.named_parameters():
    print(name, param.size())

fc1.weight torch.Size([50, 13])
fc1.bias torch.Size([50])
fc2.weight torch.Size([30, 50])
fc2.bias torch.Size([30])
fc3.weight torch.Size([1, 30])
fc3.bias torch.Size([1])

변수명을 모를 경우 위와 같이 model.named_parameters()를 통해 변수명과 변수를 동시에 불러올 수 있다.

 

# torchsummary의 summary 호출 (입력 tensor를 동일 디바이스로 설정)
summary(model, (13,), device=str(device))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Linear-1                   [-1, 50]             700
            Linear-2                   [-1, 30]           1,530
           Dropout-3                   [-1, 30]               0
            Linear-4                    [-1, 1]              31
================================================================
Total params: 2,261
Trainable params: 2,261
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.01
Estimated Total Size (MB): 0.01
----------------------------------------------------------------

위와 같이 전반적인 내용을 요약하여 출력할 수 있다.