Implement a Shallow NN with PyTorch autograd

지금까지 학습한 PyTorch autograd를 이용하여 앞에서 구현해보았던 XOR 문제를 해결하는 neural network를 더 간단하게 다시 한 번 구현해보자.

먼저 이전과 똑같이 데이터를 준비해준다!

import numpy as np
import torch # PyTorch 이용

x_seeds = np.array([(0,0), (1,0), (0,1), (1,1)])
y_seeds = np.array([0, 1, 1, 0])

N = 1000
idxs = np.random.randint(0, 4, N)

X = x_seeds[idxs]
Y = y_seeds[idxs]
X = X + np.random.normal(scale = 0.25, size=X.shape)

모델 부분이다. layer를 두개로 쌓고, tanh와 sigmoid 함수를 활성함수로 이용했다. 이전과 달라진 점은 tensor를 사용했다는 것이다!

class shallow_neural_network():
    def __init__(self, num_input_features, num_hiddens):
        self.num_input_features = num_input_features
        self.num_hiddens = num_hiddens
        
        # numpy 구현과 달리 troch tensor 사용
        self.W1 = torch.randn((num_hiddens, num_input_features), requires_grad=True)
        self.b1 = torch.randn(num_hiddens, requires_grad=True)
        self.W2 = torch.randn(num_hiddens, requires_grad=True)
        self.b2 = torch.randn(1, requires_grad=True)
        
        self.tanh = torch.nn.Tanh()
        self.sigmoid = torch.nn.Sigmoid()
        
    def predict(self, x):
        z1 = torch.matmul(self.W1, x) + self.b1
        a1 = self.tanh(z1)
        z2 = torch.matmul(self.W2, a1) + self.b2
        a2 = self.sigmoid(z2)
        return a2
 
model = shallow_neural_network(2, 3)

train 부분이다. forward, backward 과정으로 이루어진다. 이전과 달라진 점은 gradient 계산 과정을 loss.backward() 한 줄로 대체 했다는 점이다!

def train(X, Y, model, lr=0.1):
    
    m = len(X)
    cost = 0.0
    for x, y in zip(X, Y):
        
        x_torch = torch.FloatTensor(x)
        a2 = model.predict(x_torch)
        
        if y==1:
            loss = -torch.log(a2+0.0001)
        else:
            loss = -torch.log(1.0001-a2)
            
        loss.backward() # gradient 계산 
        cost+=loss.item()
        
    with torch.no_grad(): # model parameter update
        model.W1 -= lr*model.W1.grad/m
        model.b1 -= lr*model.b1.grad/m
        model.W2 -= lr*model.W2.grad/m
        model.b2 -= lr*model.b2.grad/m
        
    model.W1.requires_grad = True # 다시 model parameter tracking
    model.b1.requires_grad = True
    model.W2.requires_grad = True
    model.b2.requires_grad = True
        
    return cost/m

위 과정을 거친 후에는 이전과 똑같이 준비한 데이터와 model을 이용하여 학습을 진행하면 된다.

for epoch in range(100):
    cost = train(X, Y, model, 1.0)
    if epoch%10==0:
        print(epoch, cost)

nn.Module

PyTorch에서는 neural networks를 나타내기위해 module을 사용한다. Module은 Computation graph의 sub graph를 나타내며, autograd와 긴밀하게(?) 연결된다. 그리고 CPU와 GPU 간 저장, 불러오기 등이 쉽다!

Module을 사용하여 Train 하는 과정은 다음과 같다. module을 상속받아서 __intit__ 과 forward를 구성한다음 model.predict(x)가 아니라 바로 model(x)와 같이 model에 입력을 넣으면 알아서 forward pass를 지나게 된다!

Implementing a Shallow NN with autograd and nn.Module

마지막으로 nn.Module 까지 이용해서 다시 한번 코드를 작성해보자!

먼저 이전과 똑같이 데이터를 준비해준다! 이 때 module 사용을 위해 torch.nn을, optimizer 사용을 위해 torch.optim을 import 해주었다.

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

x_seeds = np.array([(0,0), (1,0), (0,1), (1,1)])
y_seeds = np.array([0, 1, 1, 0])

N = 1000
idxs = np.random.randint(0, 4, N)

X = x_seeds[idxs]
Y = y_seeds[idxs]
X = X + np.random.normal(scale = 0.25, size=X.shape)

모델 부분이다. 달라진 점은 nn.Module을 상속 받아 model class를 구현했다는 것이다.nn.Module을 상속 받았기 때문에 init 부분에 super().__init__()을 작성해 주었다.

class shallow_neural_network(nn.Module):
    def __init__(self, num_input_features, num_hiddens):
        super().__init__()
        self.num_input_features = num_input_features
        self.num_hiddens = num_hiddens
        
        self.linear1 = nn.Linear(num_input_features, num_hiddens)
        self.linear2 = nn.Linear(num_hiddens, 1)
        
        self.tanh = torch.nn.Tanh()
        self.sigmoid = torch.nn.Sigmoid()
        
    def forward(self, x):
        z1 = self.linear1(x)
        a1 = self.tanh(z1)
        z2 = self.linear2(a1)
        a2 = self.sigmoid(z2)
        return a2
        
model = shallow_neural_network(2, 3)

필요한 파라미터, model, optimizer, cost 함수를 선언해주었다.

num_epochs = 100
lr = 1.0
num_hiddens = 3 

model = shallow_neural_network(2, num_hiddens)
optimizer = optim.SGD(model.parameters(), lr=lr)
loss = nn.BCELoss()

train 부분이다. 이전과 달라진 점은 gradient 계산 과정뿐만 아니라 parameter update도 optimizer.step()으로 한번에 해결했다는 것이다!

for epoch in range(num_epochs):
    optimizer.zero_grad()
    
    cost = 0.0
    for x,y in zip(X, Y):
        x_torch = torch.FloatTensor(x)
        y_torch = torch.FloatTensor([y])
        
        y_hat = model(x_torch)
        loss_val = loss(y_hat, y_torch)
        cost+=loss_val
        
    cost = cost/len(X)
    cost.backward()
    optimizer.step()
    
    if epoch%10==0:
        print(epoch, cost)

test 부분이다. 학습한 modeldl xor 문제를 잘 해결하는지 확인한다.

for x,y in zip(x_seeds, y_seeds):
    print(x)
    x_torch = torch.FloatTensor(x)
    y_hat = model(x_torch)
    print(y, y_hat.item())