[PyTorch] Linear Regression(선형 회귀) - Autograd

이전 포스트에서 다루었던 Gradent Tensor에 이어 간단한 선형회귀식을 구현해본다. 

 

선형회귀는 (x, y)로 정의되는 data set에 대해 각 데이터 (x, y)와 최단 거리를 이루는 직선을 찾는 것이다. 즉, 입력 x와 출력 y 사이의 관계를 모델링한다. 

 

여기서는 단순하게 $y=wx+b$를 모델링하며, $w$는 가중치(weight), $b$는 편향(bias)이다. 선형회귀의 학습 목표는 실제 값 $y$와 예측값 $\hat{y}$간의 오차를 최소화하는 것이다. 이를 위해 손실함수로는 일반적으로 평균 제곱 오차(MSE)를 사용한다. 

$$L(w, b) = \frac{1}{n} \sum_{i=1}^n (y_i - \hat{y}_i)^2$$

이때 $y_i$는 실제값이며 $\hat{y} = wx_{i} +b$로 모델의 예측값을 의미한다. 

 

가중치를 업데이트하기 위한 방법으로는 경사하강법(Gradient descent)를 사용한다. 경사하강법은 목적 함수인 손실함수를 대상으로 Gradient를 계산하여 변수를 최적화한다. 가중치 $w$와 편향 $b$는 다음과 같이 업데이트된다. $\eta$는 학습률(learning rate)이다. 

$$w \leftarrow w - \eta \cdot \frac{\partial L}{\partial w}, \quad
b \leftarrow b - \eta \cdot \frac{\partial L}{\partial b}$$

 

이제 코드를 살펴보자. 

import torch
from matplotlib import pyplot as plt

# GPU 장치 지정
device = torch.device("cuda")

torch와 matplotlib 라이브러리에서 pyplot을 import하여 학습 결과를 그래프로 확인하도록 한다. 그리고 tensor 연산을 GPU에서 실행하기 위해 device를 설정해준다. 지금은 간단한 모델이기에 CPU에서 하는 것이 낫지만, 기능을 함께 익히기 위해 GPU로 연산을 수행하도록 한다. 

 

# 리스트 range(5)를 텐서로 변환. 1차원 -> 2차원 텐서로 변환 후 GPU로 이동
x = torch.FloatTensor(range(5)).unsqueeze(1).to(device)

입력 데이터를 생성한다. range(n)은 0부터 n-1까지 정수로 구성된 리스트를 생성하므로 range(5) = [0, 1, 2, 3, 4]이다. FloatTensor는 이를 32bit 부동소수점 tensor로 변환한다. 그 결과 1차원 tensor (5, )가 된다. 

 

unsqueeze(dim)는 dim번째 차원을 추가한다. (5, ) tensor에 1번째 차원을 추가하여 2차원 tensor (5, 1)이 된다. 그 결과 x는 다음과 같이 크기가 (5, 1)인 2차원 tensor가 된다.

$$x = \begin{bmatrix}
0 \\
1 \\
2 \\
3 \\
4
\end{bmatrix}$$

to(device)를 사용해 GPU로 tensor를 이동시킨다. 이후 이 tensor에 대한 모든 연산은 GPU에서 수행된다. 

 

# y값 생성
y = 2 * x + torch.rand(5, 1, device=device)

출력 데이터를 생성한다. y는 모델의 학습 목표값으로, $y=2x$에서 조금씩 벗어난 형태를 띄도록 구성한다. torch.rand 함수를 사용하여 [0, 1) 범위의 랜덤값을 (5, 1)크기로 생성한다. 이를 통해 약간의 노이즈를 추가함으로써 $y = 2x + noise$가 된다.

$$y = \begin{bmatrix}
0 + \text{noise} \\
2 + \text{noise} \\
4 + \text{noise} \\
6 + \text{noise} \\
8 + \text{noise}
\end{bmatrix}$$

 

# 변수의 개수 저장
num_features = x.shape[1]

num_features는 입력 데이터 x의 특성(feature)의 개수를 결정한다. 선형 회귀 모델에서 가중치 $w$의 크기는 입력 $x$의 특성의 개수에 따라 결정되므로 이를 명시적으로 확인하고 사용해야 한다. 2차원 tensor에서 각 열은 하나의 특성(feature)을 나타낸다. x는 (5, 1)이므로, 5개의 샘플(인스턴스)와 1개의 특성을 가진다.

 

x.shape는 tensor x의 크기를 반환한다. x의 크기는 (5, 1)이므로, x.shape은 torch.Size([5, 1])을 반환하며 x.shape[1]은 x의 두 번째 차원의 크기를 반환한다. 즉, x.shape[1]은 1이다. 

 

# 가중치와 편향 초기화
w = torch.randn(num_features, 1, device=device, requires_grad=True)
b = torch.randn(1, device=device, requires_grad=True)

선형 회귀 모델의 예측값은 $\hat{y}=wx+b$로 계산되므로, w는 x의 각 특성에 대해 하나의 가중치를 적용할 수 있어야 한다. w와 x의 행렬곱을 수행하기 위해서는 x의 열의 개수가 w의 행의 개수와 같아야 하며, 예측값 $\hat{y}$가 각 샘플에 대해 하나의 scalar 값으로 나와야 하므로 w의 크기는 (num_features, 1)이어야 한다. 

 

b는 $wx$에 더해지는 scalar 값으로, 예측값을 조정하는 역할을 한다. b는 각 샘플에 동일하게 적용되어야 하므로 (n, 1) 크기의 tensor wx의 n 개의 샘플에 각각 적용되어야 한다. 이를 위해 b를 (1, )의 1차원 tensor로 정의하여 브로드캐스팅을 통해 모든 샘플에 적용되도록 한다. 

 

각 모델 파라미터는 모두 requires_grad=True를 설정하여 w와 b에 대한 gradient 계산을 활성화한다.

 

# 학습률 및 최적화 도구 설정
learning_rate = 1e-3
optimizer = torch.optim.SGD([w, b], lr=learning_rate)

학습률은 0.001로 설정한다. 

 

최적화 도구로 경사하강법(Gradient Descent)를 사용해 손실 함수의 기울기(Gradient)를 기반으로 w과 b를 점진적으로 최적화한다. torch.optim.SGD에서 SGD는 Stochastic Gradient Descent(확률적 경사 하강법)의 약자이다. 확률적이라는 것은 전체 데이터셋이 아닌 미니 배치나, 샘플 단위로 경사를 계산한다는 의미이다. 이 예제에서는 배치를 사용하지 않으므로 단순 경사 하강법과 동일하다. 

 

모델 파라미터 $\theta$에 대한 업데이트는 다음과 같이 이루어진다. 

$$\theta \leftarrow \theta - \eta \cdot \nabla L(\theta)$$

$\eta$는 학습률, $\nabla L(\theta)$는 손실함수 $L$의 기울기이다. 

 

코드에서 [w, b]는 업데이트할 모델 파라미터이며 lr은 학습률이다. 

 

# 학습
loss_stack = []
for epoch in range(1001):
    optimizer.zero_grad()
    y_hat = torch.matmul(x, w) + b
    loss = torch.mean((y_hat - y) ** 2)
    loss.backward()
    optimizer.step()
    loss_stack.append(loss.item())
    print(f"Epoch {epoch}: {loss.item()}")

loss_stack은 각 epoch에서 계산되는 손실 값(loss)을 저장하기 위한 리스트이다. 이를 통해 학습 과정에서 손실 값이 어떻게 변하는지 추적하고, 학습의 수렴 여부를 확인할 수 있ek.

 

optimizer.zero_grad()는 이전 epoch에서 계산된 gradient를 초기화한다. Gradient는 PyTorch에서 누적 방식으로 관리되므로 이를 초기화해주어야한다. 

 

y_hat과 loss는 각각 예측값 $\hat{y}=wx+b$와 MSE를 계산한다.

 

손실값 loss에 대해 loss.backward()는 역전파를 수행하여 $\frac{\partial L}{\partial w}$와 $\frac{\partial L}{\partial b}$를 계산한다. 계산된 gradient는 각각 w.grad와 b.grad에 저장된다.

 

optimizer.step()은 설정한 optimizer, 즉 여기서는 SGD를 통해 w와 b를 업데이트하고, loss.item()은 손실함수의 값을 scalar 값으로 변환하여 loss_stack에 추가한다. 

 

# 예측 값 계산
with torch.no_grad():
    y_hat = torch.matmul(x, w) + b

torch.no_grad()는 gradient 계산을 비활성화한다. 학습이 아닌 예측 단계이므로, 메모리 사용량 및 연산량을 줄이기 위해서 이를 비활성화하여 예측값을 계산한다. 

 

# GPU 텐서를 CPU로 이동
x_cpu = x.cpu().numpy()  # x를 CPU로 이동 후 NumPy 배열로 변환
y_cpu = y.cpu().numpy()  # y를 CPU로 이동 후 NumPy 배열로 변환
y_hat_cpu = y_hat.cpu().numpy()  # y_hat을 CPU로 이동 후 NumPy 배열로 변환

pyplot을 통한 시각화를 위해 각 tensor를 CPU로 이동시켜 numpy 배열로 변환한다. 

 

# 시각화
plt.figure(figsize=(10, 5))

# 손실 그래프
plt.subplot(121)
plt.plot(loss_stack)
plt.title("Loss")

# 예측 값과 실제 값 비교
plt.subplot(122)
plt.plot(x_cpu, y_cpu, ".b")  # 실제 값 (Ground Truth)
plt.plot(x_cpu, y_hat_cpu, "r-")  # 예측 값 (Prediction)
plt.legend(["Ground Truth", "Prediction"])
plt.title("Prediction")

plt.show()

실행하면 다음과 같은 결과를 얻을 수 있다. 

 

손실값이 점차 감소하여 0에 수렴하고, 예측값 $\hat{y}$가 실제값 $y$에 유사한 선형 모델로 나타나는 것을 확인할 수 있다.