[pytorch] Tutorial - Build the Neural Network

References

• Official PyTorch Tutorial (link)

---

Build the Neural Network

Neural network는 데이터에 대한 연산을 수행하는 layer 또는 module로 구성됩니다. 파이토치에서 torch.nn 네임스페이스는 neural network를 구성하는데 필요한 레이어들을 제공합니다. 파이토치의 모든 레이어(모듈)들은 nn.Module의 sub class이며, neural network는 다른 레이어(모듈)들로 구성된 레이어(모듈)일 수도 있습니다.

 

이번 포스팅에서 파이토치 튜토리얼에서 설명하는 neural network에 대해 살펴보도록 하겠습니다.

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

 

Get Device for Training

아시다시피 GPU와 같은 하드웨어 가속기를 사용하여 모델을 학습시킬 수 있는데, 각자 환경에 맞는 디바이스를 다음과 같이 선택할 수 있습니다.

if torch.cuda.is_available():
    device = 'cuda'
elif torch.backends.mps.is_available():
    device = 'mps'
else:
    device = 'cpu'
print(device)

 

Define the Class

네트워크는 다양한 방법으로 구성할 수 있습니다. torch.nn.Sequential을 통해 레이어들을 쌓을 수도 있고, torch.nn.Module 클래스의 서브클래스를 작성하여, 네트워크를 구성할 수도 있습니다. 여기서는 튜토리얼에서 설명하는 nn.Module의 서브클래스를 구현하여 네트워크를 구성하는 방법에 대해 살펴봅니다. 이 클래스 내부에서는 nn.Sequential을 사용하여 레이어를 구성하고 있습니다.

여기서 입력 데이터는 FashionMNIST라고 가정하고 작성되었으며, FashionMNIST 데이터의 input shape는 (1,28,28) 입니다.

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

이렇게 생성한 NeuralNetwork의 인스턴스를 다음과 같이 생성하고, 사용할 device로 설정해줍니다.

model = NeuralNetwork().to(device)
print(model)

 

생성한 모델을 사용하려면, input data를 전달하기만 하면 됩니다. 모델에 input data를 전달하면, 클래스에서 구현한 forward 함수를 실행됩니다. 내부 백그라운드 연산에 의해서 forward 함수가 호출되기 때문에 model.forward()를 직접 호출하면 안됩니다.

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

 

Model Layers

파이토치의 torch.nn 에는 많은 레이어들을 지원하며, 위에서는 nn.Flatten, nn.ReLU, nn.Linear 를 사용하고 있습니다.

그리고, nn.Sequential을 사용하여 레이어들을 순차적으로 나열하고 있는데, nn.Sequential은 모듈(레이어)를 순차적으로 쌓는 컨테이너입니다. 이렇게 정의된 모듈에 데이터를 전달하면, 내부에 정의된 순서대로 데이터가 전달됩니다.

 

파이토치에서 제공하는 building block들은 아래 링크에서 참조하실 수 있습니다.

torch.nn — PyTorch 1.13 documentation

 

Model Parameters

Neural network의 많은 레이어들은 매개변수(ex, weight, bias)를 가지고 있으며, 이들은 학습하면서 최적화됩니다. nn.Module를 상속받으면, 모델 객체 내에서 정의된 모든 필드를 자동으로 추적하며, 모델의 모든 파라미터들을 parameters() 나 named_parameters() 메소드를 통해 액세스할 수 있습니다.

print(f"Model structure: {model}\n\n")

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")