본 블로그의 모든 글은 직접 공부하고 남기는 기록입니다.
잘못된 내용이나 오류가 있다면 언제든지 댓글 남겨주세요.
오늘은 힌튼의 Forward Forward 알고리즘을 구현해보고자 합니다. 논문을 리뷰한 지는 꽤 됐는데 바빠서 손을 못 대고 있다가 이제야 포스팅을 하게 되네요. PyTorch 공식 코드가 있어서 이를 참고해서 작성합니다.
해당 논문 리뷰를 먼저 읽고 오면 이 글을 이해하는 데 도움이 됩니다.
MNIST 데이터셋을 사용합니다. 적합한 train/test loader를 만들어줍니다.
def MNIST_loaders(train_batch_size=50000, test_batch_size=10000):
transform = Compose([
ToTensor(),
Normalize((0.1307,), (0.3081,)),
Lambda(lambda x: torch.flatten(x))])
train_loader = DataLoader(
MNIST('./data/', train=True,
download=True,
transform=transform),
batch_size=train_batch_size, shuffle=True)
test_loader = DataLoader(
MNIST('./data/', train=False,
download=True,
transform=transform),
batch_size=test_batch_size, shuffle=False)
return train_loader, test_loader
은 각각 MNIST dataset의 mean과 std입니다.
이 글에서는 지도학습일 경우만 구현합니다.
FF 알고리즘을 사용한 지도학습에는 negative sample이 필요합니다. 논문에서는 이미지의 첫 10 픽셀을 one hot encoded label로 바꿔서 사용했습니다.
def overlay_y_on_x(x, y):
x_ = x.clone()
x_[:, :10] *= 0.0 # 첫 10 픽셀을 0으로 초기화
x_[range(x.shape[0]), y] = x.max() # y 자리를 non zero 값으로 변경
return x_
즉, positive sample과 negative sample은 one hot encoded label 정보만 다르고 나머지는 같은 데이터입니다.
FF 알고리즘의 핵심 아이디어는 Pos/Neg pass 모두 feed forward network를 사용한다는 것입니다. 이 때 layer 단위로 최적화를 진행하기 때문에 먼저 각 pass에 사용할 Layer
클래스부터 구현하겠습니다.
논문에서는 식 (1)과 같이 positive sample을 판정하기 위한 SSE 형태의 목적함수를 설계했습니다. 이 글에서도 해당 수식을 사용합니다.
식 (1)은 주어진 벡터 가 positive sample일 확률을 의미합니다. 즉, positive에 최적화를 시키고 싶다면 positive 여야 하고, negative 여야 한다는 점을 기억합시다.
먼저 레이어를 구성하는 변수들을 세팅합니다.
class Layer(torch.nn.Linear):
def __init__(self, in_features, out_features, bias):
super().__init__(in_features, out_features, bias)
self.threshold = 2.0
self.epochs = 1000
self.relu = nn.ReLU()
self.optim = Adam(self.parameters(), lr=0.001)
이어서 forward 함수를 만들어줍니다. 넘겨받는 입력 벡터의 길이를 normalize 해주어야 합니다.
def forward(self, x):
x = x / x.norm(2, 1, keepdim=True) + 1e-4
x = torch.mm(x, self.weight.T) + self.bias.unsqueeze(0)
x = self.relu(x)
return x
이제 epochs만큼 이를 반복하는 train
함수를 만들어줍니다. train
함수는 positive sample과 negative sample을 모두 처리할 수 있어야 하고, epoch마다 loss를 계산해야 합니다.
def train(self, x_pos, x_neg):
for i in range(self.epochs):
g_pos = self.forward(x_pos).pow(2).mean(1)-self.threshold
g_neg = self.forward(x_neg).pow(2).mean(1)-self.threshold
각 레이어에서 사용할 loss function을 설계해 보겠습니다.
기본적으로 FF 알고리즘은 true/false를 구분하는 이진분류 문제입니다. 그런데 위에서 살펴본 것과 같이 sample data를 판정하는 값에 대한 조건이 존재합니다.
지도학습일 때는 간단하게 해결할 수 있습니다. label이 존재하기 때문이죠. 즉, true sample에 최적화되도록 하면 됩니다. 이진 분류 문제이므로 loss function으로 Binary Cross Entropy를 사용합니다.
우리는 positive sample에 최적화하는 것이 목적이므로 BCE를 softplus 함수로 대체할 수 있습니다. 따라서 식 (2)를 아래와 같이 바꿔 쓸 수 있습니다.
지금까지 설계한 loss는 input으로 하나만 받습니다. 따라서 두 데이터를 하나로 합쳐서 넣어줄 수 있도록 만들어 줍니다.
g_neg
에 마이너스를 붙이는 이유는 negative sample에서 멀어지도록 하기 위함입니다.
이제 옵티마이저를 초기화하고, 그래디언트를 계산해서, 파라미터를 업데이트하는 과정을 추가해주면 완성입니다.
def train(self, x_pos, x_neg):
for i in range(self.epochs):
g_pos = self.forward(x_pos).pow(2).mean(1)-self.threshold
g_neg = self.forward(x_neg).pow(2).mean(1)-self.threshold
loss = torch.log(1 + torch.exp(torch.cat([-g_pos, g_neg]))).mean()
self.optim.zero_grad()
loss.backward()
self.optim.step()
return self.forward(x_pos).detach(), self.forward(x_neg).detach()
이제 본격적으로 네트워크를 구현합니다. 먼저 레이어 계층을 만듭니다.
class FF(torch.nn.Module):
def __init__(self, dims):
super().__init__()
self.layers = []
for d in range(len(dims)-1):
self.layers += [Layer(dims[d], dims[d+1]).cuda()]
이어서 predict 함수를 구현합니다. 논문에서는 각 레이어의 goodness를 가지고 최적화 정도를 평가하게 되어 있습니다. 또한 지도학습이므로 레이블 개수만큼 최적화를 해주어야 합니다.
def predict(self, x):
goodness_per_label = []
for label in range(10):
h = overlay_y_on_x(x, label)
goodness = []
for layer in self.layers:
h = layer(h)
goodness += [h.pow(2).mean(1)]
goodness_per_label += [sum(goodness).unsqueeze(1)]
goodness_per_label = torch.cat(goodness_per_label, 1)
return goodness_per_label.argmax(1)
이제 train 함수를 구현합니다. positive sample과 negative sample을 받아서 최적화하는 일을 합니다.
def train(self, x_pos, x_neg):
h_pos, h_neg = x_pos, x_neg
for i, layer in enumerate(self.layers):
print('training layer', i, '...')
h_pos, h_neg = layer.train(h_pos, h_neg)
print 함수를 추가하여 self.layers
를 확인할 수 있습니다.
[Layer(
in_features=784, out_features=500, bias=True
(relu): ReLU()
),
Layer(
in_features=500, out_features=500, bias=True
(relu): ReLU()
)]
이제 모델을 학습시키고, 이를 시각화해보겠습니다.
먼저 시각화하는 함수는 다음과 같습니다.
def visualize_sample(data, name='', idx=0):
reshaped = data[idx].cpu().reshape(28, 28)
plt.figure(figsize = (4, 4))
plt.title(name)
plt.imshow(reshaped, cmap="gray")
plt.show()
이어서 모델을 학습하는 함수를 만들어 줍니다. 먼저 MNIST 데이터셋을 로드합니다.
if __name__ == "__main__":
torch.manual_seed(1234)
train_loader, test_loader = MNIST_loaders()
FF 네트워크를 구성하고, 데이터셋을 차례대로 투입합니다. 모델이 예측한 결과와 실제 데이터를 비교하여 error score를 얻습니다.
net = FF([784, 500, 500])
x, y = next(iter(train_loader))
x, y = x.cuda(), y.cuda()
x_pos = overlay_y_on_x(x, y)
rnd = torch.randperm(x.size(0))
x_neg = overlay_y_on_x(x, y[rnd])
for data, name in zip([x, x_pos, x_neg], ['orig', 'pos', 'neg']):
visualize_sample(data, name)
net.train(x_pos, x_neg)
print('train error:', 1.0 - net.predict(x).eq(y).float().mean().item())
x_te, y_te = next(iter(test_loader))
x_te, y_te = x_te.cuda(), y_te.cuda()
print('test error:', 1.0 - net.predict(x_te).eq(y_te).float().mean().item())
이상으로 구현을 마칩니다. 공식 코드 덕분에 편안하게 작업했네요.
참고문헌