Add `channel_list` to `BasicGNN`

[binkjakub]

🛠 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, and num_layers parameters as an alternative to channels_list (as in MLP)

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.

[rusty1s]

I think this sounds good, thanks for taking the lead on this one. Agree that we can just disable JumpingKnowledge in these cases.

[brovatten]

Hi. I started this issue two days ago and I think I am done with the code. But when running the tests the suites from master (without my changes) it freezes. Here is where it freezes:

test/distributed/test_dist_link_neighbor_loader.py::test_dist_link_neighbor_loader_homo[None-True-0-2] SKIPPED ('sample_from_edges' not yet implemented)
test/distributed/test_dist_link_neighbor_loader.py::test_dist_link_neighbor_loader_hetero[edge_type0-None-True-0-2] SKIPPED ('sample_from_edges' not yet implemented)
test/distributed/test_dist_neighbor_loader.py::test_dist_neighbor_loader_homo[True-0-2] Computing METIS partitioning...
Done!
ERROR:root:Coroutine task failed with error: pyg::dist_neighbor_sample() expected at most 12 argument(s) but received 13 argument(s). Declaration: pyg::dist_neighbor_sample(Tensor rowptr, Tensor col, Tensor seed, int num_neighbors, Tensor? time=None, Tensor? seed_time=None, Tensor? edge_weight=None, bool csc=False, bool replace=False, bool directed=True, bool disjoint=False, str temporal_strategy="uniform") -> (Tensor, Tensor, int[])
ERROR:root:Coroutine task failed with error: pyg::dist_neighbor_sample() expected at most 12 argument(s) but received 13 argument(s). Declaration: pyg::dist_neighbor_sample(Tensor rowptr, Tensor col, Tensor seed, int num_neighbors, Tensor? time=None, Tensor? seed_time=None, Tensor? edge_weight=None, bool csc=False, bool replace=False, bool directed=True, bool disjoint=False, str temporal_strategy="uniform") -> (Tensor, Tensor, int[])

I have installed following the contributing.md in virtualenv using:

pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv -f https://data.pyg.org/whl/torch-2.1.1+cu118.html
pip install -e ".[dev,full]"

Would appreciate some help on this to get started contributing :)

[rusty1s]

Please don't worry about this. I think it should be fixable by installing pyg-lib nightly on your end, but it doesn't matter for your PR :)