描述
🛠 Proposed Refactor
Today, all pre-defined GNN models (GCN, GAT, etc.) can only have a constant hidden size. The base class for these models, which is BasicGNN, disallows the use of varying hidden channels. However, there are many use cases when one would want to instantiate a model with varying layer sizes. For instance, BGRL method utilizes GAT or GCN architectures with decreasing hidden size in consecutive layers, e.g., [512, 256]. Also, it is common to experiment with different layer configurations, and the current implementation of GNN models in PyG forces us to write this class from scratch each time.
Suggest a potential alternative/fix
I propose to refactor the implementation of BasicGNN and make it similar to MLP, i.e.:
- add
channel_list: Optional[list[int]]parameter to the__init__method, - keep
input_channels,hidden_channels, andnum_layersparameters as an alternative tochannels_list(as inMLP)
However, unlike in MLP, there is JumpingKnowledge module in BasicGNN, which requires constant hidden size. To fix this problem, one could prohibit using channel_list with varying hidden sizes together with JumpingKnowledge by raising an exception.
Let me know if you agree to that change or have some suggestions. I can prepare PR with the proposed changes.