PyTorch로 직접 구현하는 Linear Regression

들어가며

이 글을 읽기 이전에 Linear Regression에 대한 전반적인 개념 또는 Numpy 구현에 대해 참고하고자 한다면 다음 링크를 먼저 읽고 오는 것을 추천한다.

Numpy로 직접 구현하는 Linear Regression

Iris 데이터 불러오기

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# iris 데이터 불러오기
iris = load_iris()
X = iris.data[:, [0]] # 특성 데이터에서 첫 번째 특성만 사용
y = iris.data[:, 2] # 타겟 데이터로 세 번째 특성 사용

# 학습-테스트 데이터 분리
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

데이터 Tensor 형식으로 변경

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"

# 데이터 tensor 형태로 변환하기
X_train_tensor = torch.tensor(X_train, dtype=torch.float32, device=device)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32, device=device).view(-1,1)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32, device=device)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32, device=device).view(-1,1)

print(X_train_tensor.shape, y_train_tensor.shape)
print(X_test_tensor.shape, y_test_tensor.shape)

>>> torch.Size([120, 1]) torch.Size([120, 1]) 
>>> torch.Size([30, 1]) torch.Size([30, 1])

torch DataLoader

from torch.utils.data import DataLoader, TensorDataset

# DataLoader 설정하기
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

nn.Module(기본 모델 클래스)을 상속받은 Linear Regression 모델 구현

import torch.nn as nn

# nn.Module 상속받아 LinearRegression 구현
class MultipleLinearRegression(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(MultipleLinearRegression, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

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

모델 생성

input_size = X_train_tensor.shape[1]
model = MultipleLinearRegression(input_size,1)
model.to(device)

모델 학습

import torch.optim as optim
import numpy as np

criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

model.train() # 학습을 위해 모델이 gradient를 저장하도록 설정
for epoch in range(10000):
    for X, y_true in train_loader:    
        optimizer.zero_grad() # 옵티마이저의 gradient 초기화    
        y_pred = model(X)    
        loss = criterion(y_pred, y_true) # 손실 계산 (MSE)        
        loss.backward() # gradient descent 수행 (기울기 계산)        
        optimizer.step() # SGD 방식의 최적화 진행 (가중치 업데이트)        
    if (epoch + 1) % 1000 == 0:        
        print(f"Epoch {epoch+1}, Loss: {loss.item()}")

>>> Epoch 1000, Loss: 0.7009232640266418 
>>> Epoch 2000, Loss: 0.8327165246009827 
>>> Epoch 3000, Loss: 0.8950142860412598 
>>> Epoch 4000, Loss: 0.7218224406242371 
>>> Epoch 5000, Loss: 1.0261136293411255 
>>> Epoch 6000, Loss: 0.5871046185493469
>>> Epoch 7000, Loss: 0.7505958080291748
>>> Epoch 8000, Loss: 0.9513499736785889
>>> Epoch 9000, Loss: 0.7906377911567688
>>> Epoch 10000, Loss: 0.9563422799110413

모델 평가 및 시각화

import matplotlib.pyplot as plt

model.eval()
with torch.no_grad():
    y_pred = model(X_test_tensor)
    test_loss = criterion(y_pred, y_test_tensor)
    print(f"Test Loss(MSE): {test_loss.item()}")

    # 시각화
    plt.scatter(X_test_tensor, y_test_tensor, label="Data")
    plt.plot(X_test_tensor, y_pred, color="r", label="Regression Line")
    plt.title("Iris Linear Regression")
    plt.legend()
    plt.show()

>>> Test Loss(MSE): 0.6069501638412476