Using more advanced normalization layers in models based on BasicGNN

[kmaziarz]

Models based on BasicGNN have a norm argument to pass a normalization layer. However, the normalization layers build through this are then only called on x (i.e. node features), and I can't see how they would get other arguments such as batch. Normalization layers like GraphNorm need the batch argument, otherwise defaulting to normalizing across everything (and not within each graph), which is not how GraphNorm was defined in the original paper, leading to odd behaviour during inference: a model in eval mode gives different results for a fixed input depending on other inputs in the same batch.

How can I use GraphNorm in GNN Models without having to essentially copy the internals and calling the normalization layers myself?

(Also, unless I'm misunderstanding something, to me the current treatment of normalization layers borders on being a bug, as it happily uses the provided normalization layers, just in an unexpected way)

[rusty1s]

Yes, that is correct. There is currently no easy fix to that and it is hard to provide every possible combination of GNN and normalization layers into a single module. We would need to fix this by parsing the argument list of both the GNN and the normalization layer to propagate the necessary arguments accordingly (i.e., GNN layers do not need batch while normalization layers may need it).

[erikhuck]

@kmaziarz I don't think it'll work to make the batch a field in a wrapper because a different batch could be passed into the forward method each time. I think it would need to be a method parameter rather than a class member. See my solution below and tell me what you think.

[kmaziarz]

A different batch could be passed into the forward method each time.

Yes, but each time forward on my model is called, I first overwrite the batch vector stored in the wrapper, and then I call the pyg GNN. When the GNN calls the wrapped normalization layer, the wrapper will pick up the batch vector that was stored there previously.

As I said this is a horrible hack because one has to remember to update batch before every call to the GNN.

The whole issue is that when using pygs high-level GNNs, it's not the user who is calling the normalization layer, it's the GNN implementation doing so internally, and so one cannot control what inputs are passed into there. In my code snipped above I bypass this by "storing" the batch vector as an attribute of NormWrapper.

[erikhuck]

@kmaziarz Yes I understand the main issue that you've brought up here. I actually found this discussion you created because I noticed the same problem myself as I was experimenting with different normalization layers. What do you think about the solution I proposed below (for the part of the BasicGNN forward method where the normalization layer is called)? I don't think one would have to remember to update batch before every call but would simply provide the batch argument or not depending on if the underlying normalization layers support batch or not (e.g. BatchNorm does not but InstanceNorm does. So if they choose BatchNorm for the normalization, they would not be able to pass in batch to the forward layer of the GNN (it would default to None). But if they choose a normalization layer like GraphNorm, they would be able to pass in batch.

I also agree that the GraphNorm in particular should log a warning if batch is not provided rather than silently treating the entire batch as a single graph.

[kmaziarz]

Is your suggested fix below a proposed change to pyg code itself? Currently forward in BasicGNN doesn't take batch as an argument; where does batch in your code snippet come from? Would you modify the signature of forward?

(Also, to be clear: my snippet above is not a suggested change to pyg, and rather a hack that one can do in their own codebase to get the right behaviour while continuing to use unchanged pyg. Of course if one modifies pyg then it's easy to come up with a much less confusing solution.)

[erikhuck]

@kmaziarz Yes my suggestion is for modifying pyg itself and yes the signature of BasicGNN would need to be modified to include an optional batch=None parameter. Thank you for the clarification and for providing your intermediary solution.

As for escalating to a suggested code change in pyg, what do you think of my suggestion below? Can you think of edge cases? I realized one myself: The MLP would also need an optional batch argument since the normalization layers in MLPs may also be operating on graphs.

[erikhuck]

@rusty1s How about this??

if self.norms is not None:
    norm = self.norms[i]
    if batch is not None:
        try:
            norm(x, batch=batch)
         except TypeError:
            raise ValueError(f'batch specified but underlying normalization layer, {type(norm)}, does not support a batch argument. To correct this, either configure the graph neural network with a normalization layer that supports batches or do not provide a batch argument to its forward method')
    else:
        x = norm(x)

[rusty1s]

yes, something like this. We would need to add batch=None to the forward method of BasicGNN. Instead of doing try/except, I think it would be cleaner to parse the argument list of the normalization layers via inspect, and then pass batch in whenever this argument exist. Let me know if you have interest to work on this.

[erikhuck]

@rusty1s Yes I would be happy to work on this. I will fork this repo and create a pull request. However, I realize that this will need to be done in other places besides BasicGNN e.g. MLP and aggregations.

After giving this some extra thought, l think it'd be better to, instead of just checking for batch, to follow a similar pattern as the __init__ of the MLP and BasicGNN etc. where an optional norm_kwargs parameter can provide additional arguments to the norm but instead of the norm constructor, the norm's forward call. This should default to an empty dictionary such that only x is provided. It's up to the user to choose which parameters to pass in. This would enable them to be able to pass in other arguments besides batch e.g. batch_size, p, etc. depending on which normalization type they chose. Assuming I don't notice any issues with this pattern as I get into the weeds, I will implement this pattern for BasicGNN forward method, that of MLP, and the aggregations that use norm layers.

[rusty1s]

Fixed via #8024