파이토치(Pytorch): 단순 선형 회귀 모델 구현

• 수업 내용 리마인드 및 아카이빙 목적의 업로드

---

---

 

이번 글에서는 파이토치(Pytorch)를 사용하여 단순 선형 회귀 모델을 구현하고 학습시키는 과정을 단계별로 살펴볼게요. 기본적인 데이터셋 생성부터 모델 정의, 학습, 평가, 시각화까지 전체 워크플로우를 다뤄요. 이 글을 통해 선형 회귀 모델이 어떻게 작동하는지 이해하고, 파이토치를 이용해 직접 모델을 구현해봐요.

 

1. 필요한 라이브러리 로드

먼저 실습에 필요한 라이브러리부터 설치하고 가져와야겠죠? torchinfo는 모델 구조를 확인하는 데 유용한 도구에요.

pip install torchinfo

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset, random_split
from torchinfo import summary
from sklearn.metrics import r2_score
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

 

그리고, 재현 가능성을 위해 시드를 고정하고, 우리가 사용할 장치를 설정할게요. 만약 GPU가 있다면 GPU를, 없으면 CPU를 사용할 거예요.

torch.manual_seed(10)
torch.cuda.manual_seed(10)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

 

2. 데이터 생성 및 시각화

이번 실습에서는 y=2x+3y = 2x + 3y=2x+3이라는 선형 방정식에 약간의 노이즈를 더해 데이터를 만들 거예요. 그 데이터를 가지고 회귀 모델을 학습하게 될 거예요.

X = torch.linspace(0, 10, 100).reshape(-1, 1)
y = 2 * X + 3 + torch.randn_like(X)  # y = 2x + 3 + noise

# 데이터 시각화
plt.figure(figsize=(10, 5))
plt.scatter(X.numpy(), y.numpy(), label='Generated Data', color='blue')
plt.xlabel('X')
plt.ylabel('y')
plt.title('Generated Data Visualization')
plt.legend()
plt.show()

 

3. 커스텀 데이터셋 정의 및 데이터로더 생성

파이토치에서 제공하는 Dataset 클래스를 상속받아 우리만의 데이터셋을 정의해볼게요. 그리고 DataLoader를 사용해 모델이 데이터를 쉽게 처리할 수 있도록 배치 단위로 데이터를 나누어요.

class CustomDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

custom_dataset = CustomDataset(X, y)

# 데이터셋을 train, validation, test로 나눔
train_size = int(0.7 * len(custom_dataset))
val_size = int(0.15 * len(custom_dataset))
test_size = len(custom_dataset) - train_size - val_size

train_dataset, val_dataset, test_dataset = random_split(custom_dataset, [train_size, val_size, test_size])

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

 

4. 모델 정의

이제 본격적으로 모델을 만들어 볼까요? 우선 단순 선형 회귀 모델을 정의해볼게요.

class SimpleLinearRegression(nn.Module):
    def __init__(self):
        super(SimpleLinearRegression, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

 

그리고 레이어가 추가된 조금 더 복잡한 모델도 만들어볼게요.

class FixedLayerModel(nn.Module):
    def __init__(self, input_size=1, output_size=1, hidden_units1=10, hidden_units2=10, hidden_units3=10, dropout=0.5):
        super(FixedLayerModel, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_units1)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(p=dropout)

        self.layer2 = nn.Linear(hidden_units1, hidden_units2)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(p=dropout)

        self.layer3 = nn.Linear(hidden_units2, hidden_units3)
        self.relu3 = nn.ReLU()
        self.dropout3 = nn.Dropout(p=dropout)

        self.output_layer = nn.Linear(hidden_units3, output_size)

    def forward(self, x):
        x = self.relu1(self.layer1(x))
        x = self.dropout1(x)
        x = self.relu2(self.layer2(x))
        x = self.dropout2(x)
        x = self.relu3(self.layer3(x))
        x = self.dropout3(x)
        x = self.output_layer(x)
        return x

 

5. 모델 학습

이제 학습을 위한 train 함수와 모델을 평가하기 위한 evaluate 함수를 만들어 볼게요.

def train(model, dataloader, criterion, optimizer):
    model.train()
    total_loss = 0
    for X_batch, y_batch in dataloader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        optimizer.zero_grad()
        output = model(X_batch)
        loss = criterion(output, y_batch)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * X_batch.size(0)
    return total_loss / len(dataloader.dataset)

def evaluate(model, dataloader, criterion):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for X_batch, y_batch in dataloader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            output = model(X_batch)
            loss = criterion(output, y_batch)
            total_loss += loss.item() * X_batch.size(0)
    return total_loss / len(dataloader.dataset)

 

6. 훈련 루프

이제 모델을 훈련시키고, 최적의 모델을 저장하는 과정을 진행해요.

n_epochs = 100
best_valid_loss = float('inf')
best_epoch = 0
save_path = 'best_model.pth'
train_losses = []
val_losses = []

with tqdm(total=n_epochs, desc="Training", unit="epoch", leave=True) as pbar:
    for epoch in range(1, n_epochs + 1):
        train_loss = train(model, train_loader, criterion, optimizer)
        valid_loss = evaluate(model, val_loader, criterion)

        train_losses.append(train_loss)
        val_losses.append(valid_loss)

        pbar.set_description(f"Epoch {epoch}/{n_epochs}, Train Loss: {train_loss:.4f}, Validation Loss: {valid_loss:.4f}")
        pbar.update(1)

        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            best_epoch = epoch
            torch.save({
                'epoch': best_epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': best_valid_loss,
            }, save_path)

print(f"Best model saved from epoch {best_epoch} with validation loss: {best_valid_loss:.4f}")

 

7. 모델 평가 및 예측

마지막으로, 테스트 데이터에서 모델 성능을 평가하고 새로운 데이터를 예측해보는 단계입니다.

# 테스트 데이터에서 모델 성능 평가
checkpoint = torch.load('best_model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)

test_loss = evaluate(model, test_loader, criterion)
print(f'Test Loss: {test_loss:.4f}')

# 단일 샘플 예측
sample_data = torch.tensor([7.5758], dtype=torch.float32).unsqueeze(0).to(device)
with torch.no_grad():
    prediction = model(sample_data)
print(f"Predicted value: {prediction.item()}")