7-25 PyTorch 뉴럴 네트워크 nn.Module 하위클라스 코딩

nn.Sequential에서는 instance인 model.parameters()에서 .no_grad() 조건하에서 웨이트 값들을 끄집어 내어 학습과정에서 직접 업데이트를 시키면 그 다음 학습 과정에서 업데이트 된 웨이트 값들이 model 에 입력되어 사용되었다. 하지만 nn.Sequential 과는 달리 nn.Module을 하위의 클라스로 불러 처리하고자 할 경우에는 super(TwoLayerNet, self)를 사용해서 self와 함께 TwoLayerNet의 nn.Module을 상속시켜 줄 필요가 있다. 즉 nn.Module 상속이 이루어지면 업데이트 된 웨이트 값들이 얻어지며 아울러 클라스 내에서 method 인 forward 루틴에 의해 y_pred 까지 연산을 깔끔하게 마친 상태에서 학습루틴에 리턴 시킬 수 있다. 단 학습을 위한 옵티마이저로서 optim.SGD를 사용해야 하며, loss.backward() 다음 마지막에 반드시 optimizer.step()을 실행시켜 주어야 한다.

아무래도 nn.Module을 통으로 상속이 가능한 하위 클라스로 처리하는 바람에 forward 계산까지 깔끔하게 한 묶음 처리가 가능해지므로 학습 루틴 코딩이 간단해지는 장점이 있을 수 있다. PyTorch 튜토리얼로 부터 지금까지 소개했던 몇 가지 기법들이 PyTorch CNN에서 혼용되어 사용되고 있으므로 특성들을 잘 알아둘 필요가 있다.

 

#class.super.nn

import torch
import time

start_time = time.time()

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
      
        super(TwoLayerNet, self).__init__() #python 2
        #super().__init__() #python 3
       
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
      
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

N, D_in, H, D_out = 64, 1024, 100, 10

x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

model = TwoLayerNet(D_in, H, D_out)

criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for t in range(501):
   
    y_pred = model(x)
    loss = criterion(y_pred, y)
    if t % 100 == 0:
      print(t, loss.item())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

end_time = time.time()
print( "Completed in ", end_time - start_time , " seconds")