def __init__(self, net_params):
super().__init__()
in_dim = net_params['in_dim']
hidden_dim = net_params['hidden_dim']
out_dim = net_params['out_dim']
n_classes = net_params['n_classes']
num_heads = net_params['n_heads']
in_feat_dropout = net_params['in_feat_dropout']
dropout = net_params['dropout']
n_layers = net_params['L']
self.readout = net_params['readout']
self.graph_norm = net_params['graph_norm']
self.batch_norm = net_params['batch_norm']
self.residual = net_params['residual']
self.dropout = dropout
self.embedding_h = nn.Linear(in_dim, hidden_dim * num_heads)
self.in_feat_dropout = nn.Dropout(in_feat_dropout)
self.layers = nn.ModuleList([
GATLayer(hidden_dim * num_heads, hidden_dim, num_heads, dropout,
self.graph_norm, self.batch_norm, self.residual)
for _ in range(n_layers - 1)
])
self.layers.append(
GATLayer(hidden_dim * num_heads, out_dim, 1, dropout,
self.graph_norm, self.batch_norm, self.residual))
self.MLP_layer = MLPReadout(out_dim, n_classes)
self.projection_head = projection_head(out_dim, out_dim)
def __init__(self, net_params):
super().__init__()
in_dim = net_params['in_dim']
hidden_dim = net_params['hidden_dim']
n_classes = net_params['n_classes']
dropout = net_params['dropout']
self.n_layers = net_params['L']
n_mlp_layers = net_params['n_mlp_GIN'] # GIN
learn_eps = net_params['learn_eps_GIN'] # GIN
neighbor_aggr_type = net_params['neighbor_aggr_GIN'] # GIN
readout = net_params['readout'] # this is graph_pooling_type
graph_norm = net_params['graph_norm']
batch_norm = net_params['batch_norm']
residual = net_params['residual']
# List of MLPs
self.ginlayers = torch.nn.ModuleList()
self.embedding_h = nn.Linear(in_dim, hidden_dim)
for layer in range(self.n_layers):
mlp = MLP(n_mlp_layers, hidden_dim, hidden_dim, hidden_dim)
self.ginlayers.append(GINLayer(ApplyNodeFunc(mlp), neighbor_aggr_type,
dropout, graph_norm, batch_norm, residual, 0, learn_eps))
# Linear function for graph poolings (readout) of output of each layer
# which maps the output of different layers into a prediction score
self.linears_prediction = torch.nn.ModuleList()
for layer in range(self.n_layers+1):
self.linears_prediction.append(nn.Linear(hidden_dim, n_classes))
self.projection_head = projection_head(hidden_dim, hidden_dim)
if readout == 'sum':
self.pool = SumPooling()
elif readout == 'mean':
self.pool = AvgPooling()
elif readout == 'max':
self.pool = MaxPooling()
else:
raise NotImplementedError