Edge classification with simulated graph data

[mkl-joseph]

I'm working on a project where I am trying to predict whether or not observations across multiple frames in astronomical images are associated (i.e. are the same object at different times). I've posed this as a graph problem, and created a simulated dataset where nodes are observations with features position and photometric SNR. In this graph data, there are edges connecting observations that are close in position space (via Mahalanobis Distance). The edges have a binary attribute that is 1 if the objects are truly associated, and 0 otherwise (there are multiple objects in the frames). I want to train a GNN to take as input a set of observations that are all connected via the MD filtering with initial edge attributes as 0, and output the estimated edge attributes corresponding to true associations among the objects (i.e. associated objects will have 1 for edge attribute).

I have done a lot of reading of the literature and gone through examples posted here, but I am having trouble finding the right architecture to use for this. I initially thought that the link_pred.py example would fit my needs, but I do not need to predict links between nodes, just the edge attributes. Furthermore, looking through multiple discussions I found a comment here providing an example, but in my implementation (below) I am seeing no changes in the weights after updates.

The initial motivation for this project came from this paper , where they use the SpliceCNN architecture.

Any advice/suggestions/insight would be greatly appreciated.

class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)
        self.mlp = MLP([hidden_channels * 2, out_channels])

    def forward(self, data):
        src, dst = data.edge_index
        x = self.conv1(data.x, data.edge_index).relu_()
        x = self.conv2(x, data.edge_index).relu_()
        edge_attr = torch.cat([x[src], x[dst]], dim=-1)  # Generate edge embedding
        return self.mlp(edge_attr)


model = Net(graph.num_features, 64, graph.num_edge_features)
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)
criterion = torch.nn.BCEWithLogitsLoss()

def train(data):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    # print(out)
    loss = criterion(out, data.edge_attr)
    loss.backward()
    optimizer.step()
    return float(loss)

[rusty1s]

I think your understanding is fully correct, that is: You apply a GNN to learn node embeddings, and source and destination node embeddings are combined to make a prediction for the link/edge connecting these two nodes. Your code looks fine as well. Doesn't it work?

[mkl-joseph]

So out is different than pred in this example? I thought the final output of the model was the final prediction of the edge attributes?

[mkl-joseph]

Even when doing the code above, I get as out a tensor with elements much greater than 1, when I would expect this to be within the range of (0, 1) as in my training data. Then pred is all ones on my test dataset. Could it be a result of the hidden layers? Maybe using relu() isn't correct?

[rusty1s]

out can be in range (-inf, inf) when using BCE with logits loss. You need to ensure though that you are not applying a final ReLU in your MLP though. This will prevent the model from learning to push predictions to negative values.

[mkl-joseph]

Is there a ReLU performed by default in calling MLP? Or do you mean that I shouldn't perform the ReLU on the layer that is passed into the MLP (as I have it above)?

Edit: Changed the act attribute to None but still seeing no negative values in the output of my model.

[mkl-joseph]

(Thanks again for the help!!)