[Pytorch] nn.Module을 이용해 나만의 신경망 모델 만들기

딥러닝에서 모델을 정의하는 핵심은 nn.Module 클래스이다. PyTorch에서는 모든 신경망 모델이 이 클래스를 상속받아 구현된다. nn.Module을 이해하면 나만의 신경망 구조를 자유롭게 설계할 수 있다.

---

🔹 1. nn.Module이란 무엇인가

nn.Module은 PyTorch에서 모든 신경망 계층(layer)과 모델의 기반 클래스이다. 각 모델은 이 클래스를 상속받아 __init__()에서 계층을 정의하고, forward()에서 순전파 연산을 작성한다. 자동으로 학습 가능한 파라미터를 추적하고, GPU/CPU 이동도 쉽게 관리할 수 있다.

---

🔹 2. 간단한 신경망 예제

import torch
import torch.nn as nn

# 1. 모델 클래스 정의
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 입력 3, 은닉 5, 출력 2
        self.fc1 = nn.Linear(3, 5)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(5, 2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 2. 모델 생성
model = SimpleNet()
print(model)

완전연결층(Fully Connected Layer) 2개와 ReLU 활성화 함수 1개를 사용하였다. forward()에서 순전파 과정을 정의하며, 입력 x를 통과시키면 출력이 나온다.

---

🔹 3. 모델 파라미터 확인

for name, param in model.named_parameters():
    print(name, param.shape)

각 층의 가중치(weight) 와 편향(bias) 를 확인할 수 있다. 모든 학습 가능한 파라미터는 nn.Module이 자동으로 추적한다.

---

🔹 4. 순전파와 예측

x = torch.randn(1, 3)  # 입력 샘플
output = model(x)
pred = output.argmax(1)
print("예측 클래스:", pred.item())

• output : 모델이 계산한 로짓(logit)
• argmax : 가장 높은 확률의 클래스를 선택

nn.Module을 이용하면 순전파 과정을 명확하게 관리할 수 있다.

---

🔹 5. GPU로 모델 이동

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
x = x.to(device)
output = model(x)

• model.to(device) : 모델을 GPU로 이동
• 입력 텐서도 같은 device로 이동시켜야 한다
  GPU 활용으로 학습 속도를 크게 향상할 수 있다.

---

정리

• nn.Module은 PyTorch 모델의 기반 클래스이다.
• __init__()에서 계층을 정의하고, forward()에서 순전파를 작성한다.
• 모든 학습 가능한 파라미터를 자동으로 관리한다.
• GPU로 쉽게 이동 가능하며, 순전파와 예측 과정을 명확하게 구성할 수 있다.

다음 글에서는 3편: 손실 함수와 옵티마이저를 활용한 모델 학습에 대해 다루며, 실제 학습 루프를 구현한다.

---

SEO 키워드
PyTorch, 딥러닝, nn.Module, 신경망, PyTorch 입문, 머신러닝, 순전파, GPU 연산, 딥러닝 기초, PyTorch 튜토리얼