class GNN(nn.Module): def __init__(self, args): # what to do with use_gdc super(GNN, self).__init__() self.args = args self.d = args.n_node_features num_nodes = args.n_nodes K = 3 self.conv1 = ChebConv(self.d, 16, K) self.conv2 = ChebConv(16, 16, K) self.conv3 = ChebConv(16, 1, K) # then, we have to return a results for pi and v, respectively self.activ1 = nn.Linear(16, num_nodes) self.activ2 = nn.Linear(16, 1) # define parameters self.reg_params = self.conv1.parameters() self.non_reg_params = self.conv2.parameters() def forward(self, data): # print("analyzing data") # print(data.x) # print(data.edge_index) # print(data.weight) # input() x, edges, weights = data.x, data.edge_index, data.weight weights = weights.float() # print(weights) # print(weights[0]) # print(weights[0].item()) # print(type(weights[0].item())) # input() x = self.conv1(x, edges, weights) x = F.relu(x) x = F.dropout(x) x = self.conv2(x, edges, weights) x = F.relu(x) c = self.conv3(x, edges, weights) #print(c) #input() # choice = torch.masked_select(c.squeeze(), choices) choice = F.softmax(c, dim=0).view(self.args.n_nodes) v = global_mean_pool(x, torch.zeros(data.num_nodes, dtype=torch.long).to(self.args.device)) value = self.activ2(v) # print("value: ", value.squeeze()) return choice, value.squeeze()
class ChebNet(torch.nn.Module): def __init__(self, num_features, num_classes): super(ChebNet, self).__init__() self.conv1 = ChebConv(num_features, 16, K=2) self.conv2 = ChebConv(16, num_classes, K=2) self.reg_params = self.conv1.parameters() self.non_reg_params = self.conv2.parameters() def forward(self, data): x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr x = F.relu(self.conv1(x, edge_index, edge_weight)) x = F.dropout(x, training=self.training) x = self.conv2(x, edge_index, edge_weight) return F.log_softmax(x, dim=1)
class DFAChebNet(torch.nn.Module): def __init__(self, num_features, num_classes, training_method='dfa'): super(DFAChebNet, self).__init__() self.conv1 = ChebConv(num_features, 16, K=2) self.dfa_1 = DFALayer() self.conv2 = ChebConv(16, num_classes, K=2) self.dfa = DFA(dfa_layers=[self.dfa_1], no_training=training_method != 'dfa') self.reg_params = self.conv1.parameters() self.non_reg_params = self.conv2.parameters() def forward(self, data): x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr x = F.relu(self.conv1(x, edge_index, edge_weight)) x = F.dropout(x, p=0., training=self.training) x = self.dfa_1(x) x = self.conv2(x, edge_index, edge_weight) return F.log_softmax(self.dfa(x), dim=1)