Migrating to GraphGym list of issues

🐛 Describe the bug

Over the last few days I have been migrating my PyG project to GraphGym and decided to compile a list of issues and would-be-nice's I ran into during the process. Some of these things might be my misinterpretation of the intended use or things that the team is already well aware of, in which case please disregard.

1. In torch_geometric.graphgym.loader.py the get_loader function acts as a factory for different loaders. It seems that in case of sampler == 'neighbor', it should be using a NeighborLoader , as opposed to NeighborSampler. Currently, passing dataset[0] to the constructor of NeighborSampler throws an error, since it expects an edge_index instead.
2. Having accelerator='auto' in my config results in ValueError: You selected an invalid accelerator name: accelerator='cuda:0'. Available names are: cpu, cuda, hpu, ipu, mps, tpu.
3. In my original project, I used this example to build my dataset. When I tried to copy it over as is, I ran into a problem with set_dataset_info function in torch_geometric.graphgym.loader.py. Since my dataset relies on get(self, idx) method to generate data, it does not have a self.data field, which causes an error on line 224 of torch_geometric.graphgym.loader.py, since dataset.data does not exist. I was able to do some refactoring to make it work without using get(self, idx) in my case, but this seems like a limiting factor, since the same solution might not work for bigger datasets.
4. It seems like a few modules listed in default custom_graphgym are either deprecated or do not work as intended:
   • From what I understand, train is deprecated completely after recent switch to using Pytorch Lightning. It took me a while to figure out a way to use custom train and validation steps, only to realize that it doesn't seem to be possible anymore. It would be nice to keep the ability to inject custom logic into train/validation step.
   • There are some issues of how config works, which I mentioned in #5211.
   • In the example network there is a line GNNHead = register.head_dict[cfg.dataset.task]. It'd be better to change it to cfg.gnn.head, which is essentially the same, but leads to less confusion if trying to use custom heads.
   • I'm not sure what the intention of transform is. As far as I can tell, there is no way to register them and cfg.dataset.transform is never actually used. I think it'd be nice to create them based on configs and pass them to dataset creation.
5. Would be nice: I'm sure it is something on the list of things to do, but not having documentation on what built-in configs exist and how they work is making my life very difficult, since I always have to refer to torch_geometric/graphgym/config.py and then search for usages of the config of interest through the codebase. Same thing goes for torch_geometric/graphgym/register.py for figuring out what modules I can register and how they are used.
6. Would be nice: I would suggest refactoring torch_geometric/graphgym/logger.py . It is currently responsible for both logging and computing metrics. I think it'd be better to compute metrics in a separate class (which ideally would be using torchmetrics, maybe even use this class for losses as well). As for logger, I see it has been implemented through the lightning callback, so I think it would be great to be able to specify additional custom callbacks for things like Wandb, Comet, MLflow etc. (maybe even take a step further and replace scheduling and checkpointing with lightning callbacks).
7. Would be nice: having an example of how to do inference. I had to go for some pretty hacky solutions to do it, since predict_step and predict_dataloader are missing in the lightning modules.
8. Would be nice: support specifying multiple metrics to save a checkpoint on: e.g. best precision checkpoint, best recall checkpoint etc. (if ability to do custom callbacks exists, this can just be left to the user).
9. Not sure what the exact reason is for using bash scripts for grid search, but a problem I ran into is that modifying my codebase as grid search was running was also affecting the new grid search processes.

I am happy to provide additional information on any of these issues. I am also happy to help work on these time permitting.

Environment

• PyG version: 2.1.0.post1
• PyTorch version: 1.12.1
• OS: Ubuntu 22.04
• Python version: 3.10.6
• CUDA/cuDNN version: 11.6
• How you installed PyTorch and PyG (conda, pip, source): conda for torch, pip for pyg