[PyTorch] 신경망 모델 구축 방식

nn.Sequential을 사용하는 방법

nn.Sequential을 사용하여 간단한 신경망 구조를 구성할 수 있다. 이 방식은 각 층을 순차적으로 쌓아올린다.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 10),
    nn.Softmax(dim=1)
)

 

nn.Module을 사용하는 방법

nn.Module을 사용하여 사용자 정의 신경망 구조를 구성할 수 있다. 이 방식은 객체 지향 프로그래밍을 사용하여 모델을 정의한다.

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.dense1 = nn.Linear(784, 128)
        self.dense2 = nn.Linear(128, 64)
        self.dense3 = nn.Linear(64, 10)

    def forward(self, x):
        x = nn.functional.relu(self.dense1(x))
        x = nn.functional.relu(self.dense2(x))
        return nn.functional.softmax(self.dense3(x), dim=1)

model = MyModel()

 

모델을 하위 클래스로 정의하는 방법 (Model Subclassing)

이 방법은 nn.Module을 사용하여 사용자 정의 신경망 구조를 구성하는 것과 유사하다. 다만, 이 방식은 기존의 모델을 상속받아 새로운 모델을 정의하는 경우에 사용한다.

import torch
import torch.nn as nn

class BaseModel(nn.Module):
    def __init__(self):
        super(BaseModel, self).__init__()
        self.dense1 = nn.Linear(784, 128)
        self.dense2 = nn.Linear(128, 64)

    def forward(self, x):
        x = nn.functional.relu(self.dense1(x))
        return nn.functional.relu(self.dense2(x))

class MyModel(BaseModel):
    def __init__(self):
        super(MyModel, self).__init__()
        self.dense3 = nn.Linear(64, 10)

    def forward(self, x):
        x = super().forward(x)
        return nn.functional.softmax(self.dense3(x), dim=1)

model = MyModel()