I am trying to implement a TopKPooling layer but the output is not documented. I am get a tuple of 6 objects – the first two items look to be the new node feature matrix and edge indices, the third is a None type (probably because I am doing standard TopKPooling), the fourth looks to be the new batch tensor which indicates which of the nodes belong to which graph. The fifth and sixth I am less sure on: they are both 1D tensors of length |V’| where V’ is the new coarsened node set – the fifth looks to be strictly integers with no gradient attached so it could be some kind of labelling and the sixth looks to be floats with a gradient.
could some please help me understand what these outputs are, thanks
The code I am using is as follows:
class KPooling(nn.Module):
def __init__(self, ratio, input_size, hidden_size, output_size):
super(KPooling, self).__init__()
self.gcn1 = SAGEConv(input_size, hidden_size)
self.first_pool_layer = TopKPooling(hidden_size, ratio)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = self.gcn1(x, edge_index)
x, edge_index, _, batch, ?, ? = self.first_pool_layer(x, edge_index, batch=batch)
where the ?
I’ve added to show that I am unsure on what these two tensor outputs are.