[pytorch] MLP model source code by python

 

아래는 pytorch로 작성한 MLP모델의 소스코드이다.

class MLPModel(nn.Module):
  def __init__(self, in_dim, out_dim, hid_dim, n_layer, act):
    super(MLPModel, self).__init__()
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.hid_dim = hid_dim
    self.n_layer = n_layer
    self.act = act

    self.fc = nn.Linear(self.in_dim, self.hid_dim)
    self.linears = nn.ModuleList()

    for i in range(self.n_layer-1):
      self.linears.append(nn.Linear(self.hid_dim, self.hid_dim))
    
    self.fc2 = nn.Linear(self.hid_dim, self.out_dim)

    if self.act == 'relu':
      self.act = nn.ReLU()

  def forward(self, x):
    x = self.act(self.fc(x))
    for fc in self.linears:
      x = self.act(fc(x))
    x = self.fc2(x)
    return x

net = MLPModel(3072, 10, 100, 4, 'relu')

 

주요 함수

nn.Linear() - Fully connected layer function

nn.ReLU() - Activation function

nn.ModuleList() - 여기에 레이어들을 넣어야 optimizer에 등록이 된다.

                            단순 List()에 넣게되면, 히든 레이어 수를 바꿔도 파라미터 수가 같다. 또는 파라미터를 바꿔도 accuracy가 같다.