[Python/PyTorch] Sequential()을 forward()에 사용하기

문제

class GNNs(torch.nn.Module):
  def __init__(self, dataset):
    super().__init__()
    self.data = dataset

    self.conv1 = GCNConv(dataset.num_node_features, 32)
    self.conv2 = GCNConv(32, dataset.num_classes)

  def forward(self, data):
    x, edge_index = data.x, data.edge_index

    x = self.conv1(x, edge_index)
    x = F.leaky_relu(x)
    x = self.conv2(x, edge_index)

    return F.softmax(x, dim=1)


cora_model = GNNs(cora_dataset)
cora_optim = torch.optim.Adam(cora_model.parameters())

cora_model.train()
for epoch in range(500):
    cora_optim.zero_grad()
    out = cora_model(cora_data)
    loss = F.cross_entropy(out[cora_data.train_mask], cora_data.y[cora_data.train_mask])
    loss.backward()
    cora_optim.step()

cora_model.eval()
_, pred = cora_model(cora_data).max(dim=1)
correct = float (pred[cora_data.test_mask].eq(cora_data.y[cora_data.test_mask]).sum().item())
acc = correct / cora_data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))

>>> Accuracy: 0.7780

cora_model = Sequential('x, edge_index', [
            (GCNConv(cora_dataset.num_node_features, 32), 'x, edge_index -> x'),
            LeakyReLU(),
            (GCNConv(32, cora_dataset.num_classes), 'x, edge_index -> x'),
            Softmax(dim=1)
        ])

cora_model.train()
for epoch in range(500):
    cora_optim.zero_grad()
    out = cora_model(cora_data.x, cora_data.edge_index)
    loss = F.cross_entropy(out[cora_data.train_mask], cora_data.y[cora_data.train_mask])
    loss.backward()
    cora_optim.step()

cora_model.eval()
_, pred = cora_model(cora_data.x, cora_data.edge_index).max(dim=1)
correct = float (pred[cora_data.test_mask].eq(cora_data.y[cora_data.test_mask]).sum().item())
acc = correct / cora_data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))

>>> Accuracy: 0.1500

Sequential을 사용해서 동일한 레이어층의 모델을 만들었는데 정확도가 너무 다르게 나옴

해결

class GNNs(torch.nn.Module):
  def __init__(self, dataset):
    super().__init__()
    self.data = dataset

    self.model = Sequential('x, edge_index', [
            (GCNConv(cora_dataset.num_node_features, 32), 'x, edge_index -> x'),
            LeakyReLU(),
            (GCNConv(32, cora_dataset.num_classes), 'x, edge_index -> x'),
            Softmax(dim=1)
        ])

  def forward(self, data):
    x, edge_index = data.x, data.edge_index

    x = self.model(x, edge_index)

    return x


cora_model = GNNs(cora_dataset)
cora_optim = torch.optim.Adam(cora_model.parameters())

cora_model.train()
for epoch in range(500):
    cora_optim.zero_grad()
    out = cora_model(cora_data)
    loss = F.cross_entropy(out[cora_data.train_mask], cora_data.y[cora_data.train_mask])
    loss.backward()
    cora_optim.step()

cora_model.eval()
_, pred = cora_model(cora_data).max(dim=1)
correct = float (pred[cora_data.test_mask].eq(cora_data.y[cora_data.test_mask]).sum().item())
acc = correct / cora_data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))

Sequential을 써도 forward()에서 설정해주지 않아 훈련을 할 때 제대로 반영되지 않은 듯...