Novozymes 효소 안정성 예측 경진대회 우승자 Eggplanck 팀의 솔루션 발표

import dgl
import torch
import torch.nn as nn
from dgl.nn import GINConv

class GINModel(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.conv1 = GINConv(nn.Linear(in_dim, hidden_dim), 'sum')
        self.conv2 = GINConv(nn.Linear(hidden_dim, hidden_dim), 'sum')
        self.conv3 = GINConv(nn.Linear(hidden_dim, hidden_dim), 'sum')
        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim * 3, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, out_dim)
        )

    def forward(self, g, h):
        h1 = torch.relu(self.conv1(g, h))
        h2 = torch.relu(self.conv2(g, h1))
        h3 = torch.relu(self.conv3(g, h2))
        # Concatenate intermediate representations
        h_combined = torch.cat([h1, h2, h3], dim=1)
        g.ndata['h'] = h_combined
        return self.mlp(dgl.readout_nodes(g, 'h', op='sum'))