Question about passing node features

[ray-wynne]

I'm implementing a form of EdgeConv for a physics application that requires me to pass in a feature of each neighboring node. If each node j has a feature a_j, then for node i, I will aggregate over my neighbors as x_i' = sum_j a_j * h(x_i, x_j) I scale each pairwise activation by the neighbor nodes feature. I also don't pass the a_j feature to the layers.

Would I be able to make a distinction between the features for learning and the a_j features with Pytorch's Dataloaders? Or should I try to store these a_j separately?

[rusty1s]

It seems that your a_j tensor is just a scalar for each edge, right? You can simply pass that to the module as well, also I would rename it to something like edge_weight to make it more clear that this tensor holds a value per edge:

class EdgeConv(MessagePassing):
    def __init__(self, F_in, F_out):
        super(EdgeConv, self).__init__(aggr='max')  # "Max" aggregation.
        self.mlp = Seq(Lin(2 * F_in, F_out), ReLU(), Lin(F_out, F_out))

    def forward(self, x, edge_index, edge_weight):
        # x has shape [N, F_in]
        # edge_index has shape [2, E]
        return self.propagate(edge_index, x=x, edge_weight=edge_weight)  # shape [N, F_out]

    def message(self, x_i, x_j, edge_weight):
        # x_i has shape [E, F_in]
        # x_j has shape [E, F_in]
        edge_features = edge_weight * torch.cat([x_i, x_j - x_i], dim=1)  # shape [E, 2 * F_in]
        return self.mlp(edge_features)  # shape [E, F_out]

[ray-wynne]

For edge_weight, should I store this under edge_attribute for each graph? That's what I'm doing now and it seems to be working, but I'm not sure if there's a better way.

[rusty1s]

You can store it under any name you like. We typically use edge_weight to denote scalar edge weights, and edge_attr to denote multi-dimensional edge features.