Different embeddings for different node types

[douglasrizzo]

❓ Questions & Help

Hi there. I am working in a reinforcement learning problem where I model the state of my environment as a graph in which nodes may have different types, with features vectors that may have either different lengths or different meaning.
In order to aggregate messages of nodes with different types, I was wondering if it would be possible to apply different embedding networks for each node type, ensuring that the resulting embeddings have the same dimensions for all node types and that different node types are embedded by NNs specialized in each node type. I have seen something similar mentioned in #944 and specially #598, but wanted to be reassured before going full-on dev mode.

I would also like to know if you know of anyone that has used PyG in an RL task, as the training process of NNs is usually different, involving the creation of a replay buffer before training begins, as well as the use of different loss functions. I have no idea what to expect. Thanks in advance.

[rusty1s]

This is indeed possible, and there are two possible ways to achieve this.

1. Make use of something like RGCNConv which will make aggregation dependent on the edge type. This way, all node embeddings of different types will get updated simultaneously.
2. Explicitly define which node type should get updated. Here you need to maintain node and edge features of different types separately, but are free in choosing the operator and the dimensionality, e.g.:

x_A = conv1(x_A, edge_index_A)
x_B = conv2((x_A, xB), edge_index_A_B)
x_B = conv1(x_B, edge_index_B)

where you have two node types A and B, and want to first update node features of type A, send them to B, and update node features of type B afterwards.

[douglasrizzo]

Thanks for the reply. I'll try and start simple. With RGCN, I create a single graph data, in which data.x is a Tensor where all nodes have feature vectors of equal length, correct? And RGCNConv tells nodes apart through the edge_type parameter.

So, if I have nodes of type a, b and c, the way I treat them differently is by assigning the edge_type between nodes a and a to 0, a and b to 1, a and c to 2 and so on?

[rusty1s]

Exactly :)

[douglasrizzo]

I am having a specific problem and was hoping I could get some clarification. Since my data has different dimensions for each node type, how should I store it in the torch_geometric.data.Data object?

First, I went with a list of tensors (one for each node) but I found out that's not good for parallelization and had trouble creating Batch objects afterwards.

My other idea is to store everything in a single tensor and sort which lines belong to each node type inside the neural network.

[rusty1s]

What about padding all nodes features to equal size, and use RGCN as usual?

[douglasrizzo]

I padded all node features to equal size and removed the padding once inside the forward() method. I really want to test the encoding layer, even if zero padding produces similar results.

def forward(self, data):
    x, edge_index, node_type = data.x, data.edge_index, data.node_type

    # create tensor to hold the encoded results
    X = torch.zeros(x.shape[0], self.encoding_size, device=self.device)

    for nt in node_type.unique():
        encoding_layer = self.encodings[nt]  # grab torch.nn.Linear layer for node type
        node_indices = torch.nonzero(node_type == nt).squeeze()  # grab nodes of that type
        x_sub = x[node_indices, :encoding_layer.in_features
                    ]  # grab features only of those nodes, remove padding
        enc = encoding_layer(x_sub)  # apply layer to input
        X[node_indices] = enc  # put outputs in their corresponding places

I think I'll close this as I feel this problem has been solved. Thank you very much for your help @rusty1s .