Problem using HeteroBatchNorm (Example missing)

[Rico2000]

I'm having trouble using HeteroBatchNorm I don't quite understand how the forward function works and in particular the type_vec parameter how can I create it? Is there an example for the use of HeteroBatchNorm and HeteroLayerNorm my code looks like this:

class GNN(torch.nn.Module):
    def __init__(self, channels,metadata):
        super().__init__()
        self.layers = torch.nn.ModuleList([HGTConv(-1, 64,metadata),
                                           HeteroBatchNorm(64,num_types=5),
                                           HGTConv(64, 32,metadata),
                                           HeteroBatchNorm(32,num_types=5),
                                           HGTConv(32, 16,metadata),
                                           HeteroBatchNorm(16,num_types=5) ])
     self.acivate = torch.nn.LeakyReLU()

    def forward(self, x, edge_index):
        x = self.layers[0](x, edge_index)
        x = self.layers[1](x)
        x = self.acivate(x)
        x = self.layers[2](x, edge_index)
        x = self.layers[3](x)
        x = self.acivate(x)
        x = self.layers[4](x, edge_index)
        x = self.layers[5](x)
        x = self.acivate(x)
        return x

Thanks

[rusty1s]

HeteroBatchNorm operates on a single feature matrix, and type_vec denotes the node type for each row of this feature matrix, e.g.:

x = ... [num_nodes, num_features]
type_vec = [0, 0, 1, 1, 1, ...]  # First two nodes belong to node type 0, etc

In your case, the output is a dictionary, so something like

torch.nn.ModuleDict({node_type: BatchNorm(64) for node_type in node_types})
x_dict = {node_type: norm_dict[node_type](x) for node_type, x, in x_dict.items()}

already does the job.

[Rico2000]

okay thanks in which context do i usw HeteroBatchNorm insted of the example you shown?

[rusty1s]

You can use it whenever you don't want to utilize for-loops and do message passing all at once. This can have some advantages when number of node types or number of edge types get really big. For example, you can use RGCNConv together with HeteroBatchNorm.