class GCN(nn.Module):
def __init__(self, nfeat, nhid, nclass, dropout):
super(GCN, self).__init__()
self.gc1 = GraphConvolution(nfeat, nhid)
self.gc2 = GraphConvolution(nhid, nhid)
self.gc3 = GraphConvolution(nhid, nclass)
self.dropout = dropout
def forward(self, x, adj, eval=False):
x = F.relu(self.gc1(x, adj))
x = F.dropout(x, self.dropout, training=self.training)
x = self.gc2(x, adj)
if eval:
return x
else:
x = F.dropout(x, self.dropout, training=self.training)
x = self.gc3(x, adj)
return F.log_softmax(x, dim=1)
def functional_forward(self, x, adj, weights, eval=False):
x = F.relu(self.gc1.functional_forward(x, adj, id=1, weights=weights))
x = F.dropout(x, self.dropout, training=self.training)
x = self.gc2.functional_forward(x, adj, id=2, weights=weights)
if eval:
return x
else:
x = F.dropout(x, self.dropout, training=self.training)
x = self.gc3.functional_forward(x, adj, id=3, weights=weights)
return F.log_softmax(x, dim=1)
class GCN(nn.Module):
def __init__(self, nfeat, nhid, nclass, dropout):
super(GCN, self).__init__()
self.gc1 = GraphConvolution(nfeat, nhid)
self.gc2 = GraphConvolution(nhid, nhid)
self.mlp = nn.Linear(nhid, nclass)
self.nhid = nhid
self.dropout = dropout
def forward(self, graph_list, node_feat, edge_feat, eval=False):
graph_sizes, n2n_sp, e2n_sp, subg_sp, node_degs = process_sparse(graph_list, node_feat, edge_feat)
x = F.relu(self.gc1(node_feat, n2n_sp, node_degs), inplace=True)
x = F.dropout(x, self.dropout, training=self.training)
x = self.gc2(x, n2n_sp, node_degs)
batch_graphs_repr = torch.zeros(len(graph_sizes), self.nhid).cuda()
batch_graphs_repr = Variable(batch_graphs_repr)
accum_count = 0
if eval:
for i in range(subg_sp.size()[0]):
batch_graphs_repr[i, :] = torch.mean(x[accum_count: accum_count + graph_sizes[i]], dim=0, keepdim=True)
accum_count += graph_sizes[i]
return batch_graphs_repr
else:
x = F.dropout(x, self.dropout, training=self.training)
for i in range(subg_sp.size()[0]):
batch_graphs_repr[i, :] = torch.mean(x[accum_count: accum_count + graph_sizes[i]], dim=0, keepdim=True)
accum_count += graph_sizes[i]
batch_graphs_repr = self.mlp(batch_graphs_repr)
return F.log_softmax(batch_graphs_repr, dim=1)
def functional_forward(self, graph_list, node_feat, edge_feat, weights, eval=False):
graph_sizes, n2n_sp, e2n_sp, subg_sp, node_degs = process_sparse(graph_list, node_feat, edge_feat)
x = F.relu(self.gc1.functional_forward(node_feat, n2n_sp, node_degs, id=1, weights=weights), inplace=True)
x = F.dropout(x, self.dropout, training=self.training)
x = self.gc2.functional_forward(x, n2n_sp, node_degs, id=2, weights=weights)
batch_graphs_repr = torch.zeros(len(graph_sizes), self.nhid).cuda()
batch_graphs_repr = Variable(batch_graphs_repr)
accum_count = 0
if eval:
for i in range(subg_sp.size()[0]):
batch_graphs_repr[i, :] = torch.mean(x[accum_count: accum_count + graph_sizes[i]], dim=0, keepdim=True)
accum_count += graph_sizes[i]
return batch_graphs_repr
else:
x = F.dropout(x, self.dropout, training=self.training)
for i in range(subg_sp.size()[0]):
batch_graphs_repr[i, :] = torch.mean(x[accum_count: accum_count + graph_sizes[i]], dim=0, keepdim=True)
accum_count += graph_sizes[i]
batch_graphs_repr = torch.mm(batch_graphs_repr, weights['mlp.weight'].t()) + weights['mlp.bias']
return F.log_softmax(batch_graphs_repr, dim=1)
