윤제로의 제로베이스

nn.Module로 구현하는 선형회귀 본문

Background/Pytorch 기초

nn.Module로 구현하는 선형회귀

윤_제로 2022. 1. 14. 00:44

1. 단순 선형 회귀 구현

 

import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(1)

# 데이터
x_train = torch.FloatTensor([[1], [2], [3]])
y_train = torch.FloatTensor([[2], [4], [6]])

 

nn.Linear()는 입력의 차원, 출력의 차원을 인수로 받는다.

# 모델을 선언 및 초기화. 단순 선형 회귀이므로 input_dim=1, output_dim=1.
model = nn.Linear(1,1)

 

model.parameters()함수에 model()의 가중치 W와 편향 b가 저장되어 있다.

이때 W와 b는 랜덤으로 초기화 된다.

또한 두 값 모두 학습 대상이므로 requires_grad=True로 되어 있다.

 

옵티마이저를 정의할 떄 model.parameters()를 사용하여 W와 b를 전달한다.

# optimizer 설정. 경사 하강법 SGD를 사용하고 learning rate를 의미하는 lr은 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) 

# 전체 훈련 데이터에 대해 경사 하강법을 2,000회 반복
nb_epochs = 2000
for epoch in range(nb_epochs+1):

    # H(x) 계산
    prediction = model(x_train)

    # cost 계산
    cost = F.mse_loss(prediction, y_train) # <== 파이토치에서 제공하는 평균 제곱 오차 함수

    # cost로 H(x) 개선하는 부분
    # gradient를 0으로 초기화
    optimizer.zero_grad()
    # 비용 함수를 미분하여 gradient 계산
    cost.backward() # backward 연산
    # W와 b를 업데이트
    optimizer.step()

    if epoch % 100 == 0:
    # 100번마다 로그 출력
      print('Epoch {:4d}/{} Cost: {:.6f}'.format(
          epoch, nb_epochs, cost.item()
      ))

  • H 식에 입력 x로 부터 예측된 y를 얻는 것을 forward 연산이라고 한다.
  • 학습 전, prediction = model(x_train)은 x_train으로부터 예측값을 리턴하므로 forward 연산이다.
  • 학습 후, pred_y = model(new_var)는 임의의 값 new_var로부터 예측값을 리턴하므로 forward 연산이다.
  • 학습 과정에서 비용 함수를 미분하여 기울기를 구하는 것을 backward 연산이라 한다.
  • cost.backward()는 비용 함수로부터 기울기를 구하라는 의미이며 backward 연산이다.

2. 다중 선형 회귀 구현

 

import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(1)

# 데이터
x_train = torch.FloatTensor([[73, 80, 75],
                             [93, 88, 93],
                             [89, 91, 90],
                             [96, 98, 100],
                             [73, 66, 70]])
y_train = torch.FloatTensor([[152], [185], [180], [196], [142]])
# 모델을 선언 및 초기화. 다중 선형 회귀이므로 input_dim=3, output_dim=1.
model = nn.Linear(3,1)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)
nb_epochs = 2000
for epoch in range(nb_epochs+1):

    # H(x) 계산
    prediction = model(x_train)
    # model(x_train)은 model.forward(x_train)와 동일함.

    # cost 계산
    cost = F.mse_loss(prediction, y_train) # <== 파이토치에서 제공하는 평균 제곱 오차 함수

    # cost로 H(x) 개선하는 부분
    # gradient를 0으로 초기화
    optimizer.zero_grad()
    # 비용 함수를 미분하여 gradient 계산
    cost.backward()
    # W와 b를 업데이트
    optimizer.step()

    if epoch % 100 == 0:
    # 100번마다 로그 출력
      print('Epoch {:4d}/{} Cost: {:.6f}'.format(
          epoch, nb_epochs, cost.item()
      ))