Various updates to graphium #449

WenkelF · 2023-08-28T13:50:21Z

Changelogs

Further improvement of caching logic (only calling datamodule._make_multitask_dataset() once)
Integrated LargeMix into hydra
Integrated single datasets (LargeMix) into hydra
Computing val/test metrics on cpu to save gpu memory
Implemented testing using a model checkpoint

Remarks

Regarding(3.): Single dataset runs are specific to NeurIPS experiment, where we use the same model as for LargeMix, i.e., run them like follows:

graphium-train \
    architecture=largemix \
    tasks=l1000_mcf7 \
    training=largemix \

Regarding (5): Needs to be properly changed to typer following comment by @cwognum (incl. changing documentation) once branch has been updated with main.

Checklist:

[ ] Was this PR discussed in an issue? It is recommended to first discuss a new feature into a GitHub issue before opening a PR.
Add tests to cover the fixed bug(s) or the new introduced feature(s) (if appropriate).
Update the API documentation is a new function is added, or an existing one is deleted.
Write concise and explanatory changelogs above.
If possible, assign one of the following labels to the PR: feature, fix or test (or ask a maintainer to do it for you).

codecov · 2023-08-28T13:56:50Z

Codecov Report

Merging #449 (3cf2fb5) into main (3aa098b) will increase coverage by 0.01%.
Report is 3 commits behind head on main.
The diff coverage is 66.66%.

@@            Coverage Diff             @@
##             main     #449      +/-   ##
==========================================
+ Coverage   64.60%   64.62%   +0.01%     
==========================================
  Files          93       93              
  Lines        8417     8423       +6     
==========================================
+ Hits         5438     5443       +5     
- Misses       2979     2980       +1

Flag	Coverage Δ
unittests	`64.62% <66.66%> (+0.01%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Components	Coverage Δ
ipu	`49.14% <ø> (ø)`

DomInvivo · 2023-08-28T14:26:45Z

@maciej-sypetkowski , do you think that moving the val/test metrics computations to CPU would significantly hurt parallelization? We are currently constrained as we don't use the torchmetrics ability to update stats due to their inability to handle NaNs.

DomInvivo · 2023-08-29T21:10:17Z

expts/hydra-configs/architecture/largemix.yaml

All those yaml files seem like duplicate from the files in the neurips2023 folder. What do you think? Is there a reason to have them separately?

I curated the configs to make LargeMix configs easy to use within our hydra logic. In particular, the modularity with the added files will facilitate finetuning on pretrained models on LargeMix. As far as I see, neurips2023 configs can only be used for training on LargeMix, but one needs to write new configs for, e.g., finetuning or preparing data.

DomInvivo · 2023-08-29T21:10:55Z

graphium/config/_loader.py

    if "model_checkpoint" in cfg_trainer.keys():
-        dirpath = cfg_trainer["model_checkpoint"]["dirpath"]
-        path = fs.join(dirpath, path)
+        dirpath = cfg_trainer["model_checkpoint"]["dirpath"] + str(cfg_trainer["seed"]) + "/"


Can you create a unique identifier that you can associate with a specific run? Perhaps take the identifier from Wandb?

Unless you want to add an ID that is only programmatically accessible, I would leave it to the config and not hard-code the seed in the path since it is pretty subjective where you want to save the models. Note that setting the dirpath to include the seed can currently also be achieved by changing the config - either directly or through the CLI.

We have the issue that we want identifiable checkpoints, and non- overwriting checkpoints. How do you suggest we move forward?

I understand the use-case!

As I said before, I suggest using the config. I would set a sensible default, but then leave it up to the user.

E.g. the line that is hard coded here can already be configured with hydras override syntax, using OmegaConf's interpolation:

graphium-train ... "trainer.model_checkpoint.dirpath='/path/to/checkpoints/${constants.seed}'"

In case of a sweep, you might want to use something more specific, e.g.:

graphium-train -m ... "trainer.model_checkpoint.dirpath='/path/to/checkpoints/${now:%Y-%m-%d_%H-%M-%S}'"

To prevent most users from having to set a value themselves, you could include a sensible default in the config files. Current defaults (e.g. models_checkpoints/neurips2023-small-gcn/) are too specific and not consistent. I would e.g. save to the current working dir by default:

trainer: model_checkpoint: dirpath: ${hydra.run.dir}

Thanks for the suggestion. I'm still getting used to hydra and didn't know you could put the date-time directly in the Yaml! What do you think of this commit? Making it the default in all the hydra-configs

I agree that we have to leave the checkpoint path configuration somewhat up to the user.

However, the seed is an important culprit leading to overwriting (or unclear identity) of model checkpoints. I post-pended it to the dirpath because we also used to have the following in load_trainer:

if "model_checkpoint" in cfg_trainer.keys(): cfg_trainer["model_checkpoint"]["dirpath"] += str(cfg_trainer["seed"]) + "/" callbacks.append(ModelCheckpoint(**cfg_trainer["model_checkpoint"]))

I am using graphium-train ... "trainer.model_checkpoint.dirpath='/path/to/checkpoints/${constants.seed}'" as proposed by Cas to reasonbly structure checkpoints

Why not use the full date and time? Much better than the random seed. And avoids collision when doing sweeps with pre-selected seeds. And also interpretable.

Yes, I agree. It was necessary at the time the branch was created because the seed was hard-coded into the dirpath back then... I think it is good as of now using the date :)

Changed function discussed here to its current state in the main branch. The function was outdated in this branch because we added it to #441 where it was modified to accommodate further changes :)

DomInvivo · 2023-08-29T21:15:38Z

tests/test_datamodule.py

@@ -55,8 +55,7 @@ def test_ogb_datamodule(self):
            rm(TEMP_CACHE_DATA_PATH, recursive=True)

        # Reset the datamodule
-        ds._data_is_prepared = False
-        ds._data_is_cached = False
+        ds = GraphOGBDataModule(task_specific_args, **dm_args)


What's the purpose of this??

This is properly resetting the datamodule. The old code didn't reset self.train_ds, etc., which could lead to wrong testing as, e.g., setup checks if self.train_ds is None

maciej-sypetkowski · 2023-08-31T13:50:33Z

@maciej-sypetkowski , do you think that moving the val/test metrics computations to CPU would significantly hurt parallelization?

I think it should be fine, however the predictions are still kept on gpus if I look correctly (the casting is only in _general_epoch_end).

We are currently constrained as we don't use the torchmetrics ability to update stats due to their inability to handle NaNs.

It can be done in torchmetrics, you just need to implement a wrapper for ignoring nans. We'll probably want to do this when we add better multi-device support.

DomInvivo · 2023-08-31T14:14:05Z

@maciej-sypetkowski

It can be done in torchmetrics, you just need to implement a wrapper for ignoring nans. We'll probably want to do this when we add better multi-device support.

We already have the MetricWrapper, but it works on the metric level, not on the update and compute level. I guess that's a good place to start looking in the future.

…_test.py

WenkelF · 2023-08-31T23:10:33Z

@DomInvivo I think we can merge. I suggest opening a new branch for the optimizations discussed with @maciej-sypetkowski

wenkelf and others added 26 commits August 18, 2023 20:27

Adding largemix configs

99ffea5

Further improving caching

c59bd63

Computing val/test metrics on cpu to save gpu memory

290e3da

Implementing testing for a model checkpoint

ea9ff5d

Adding single dataset configs for LargeMix to hydra

fd6e932

Adding script for test sweep

77a28a0

scripts for sbatch

af948e1

correction

7b0fdf3

Switching to V100

17a92fe

Train script

00cedf9

Changing back to cudatoolkit in env.yml

fe4ead7

Updating test sweep

d6f4ae0

Updating test sweep

2a4a129

Adding single run scripts

9ba7cae

Updating pdba runs

b5d897e

Updating pcba runs

952a145

Changing split_names to test-seen

38e7667

Updating split_paths to test_seen

914cb08

Updating test_splits to test_seen

bd3784c

Removing scripts

2ed5f1e

Minor reorganization

eb037de

Enabling to resume training

cf88c3f

Finalizing largemix and large single dataset configs for hydra

5781dcb

Temporary removing code for resuming training in favor of dedicated pr

93d12be

Cleaning up

8a80c52

Reformatting with black

67e4134

WenkelF added fix PR that fixes an issue feature labels Aug 28, 2023

WenkelF requested a review from DomInvivo as a code owner August 28, 2023 13:50

DomInvivo linked an issue Aug 29, 2023 that may be closed by this pull request

Multi-level fine-tuning and "fingerprinting" #347

Closed

DomInvivo removed a link to an issue Aug 29, 2023

Multi-level fine-tuning and "fingerprinting" #347

Closed

Merge branch 'main' into caching

53d10aa

DomInvivo reviewed Aug 29, 2023

View reviewed changes

Added the date-time in the model_checkpoint paths

da31797

WenkelF added 3 commits August 31, 2023 18:45

Removing graphium/cli/test.py in favor of graphium/cli/train_finetune…

44d241f

…_test.py

Minor fix

8f1ddfb

Updating get_checkpoint_path

3cf2fb5

DomInvivo merged commit 0d24634 into main Sep 1, 2023
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Various updates to graphium #449

Various updates to graphium #449

WenkelF commented Aug 28, 2023 •

edited

Loading

codecov bot commented Aug 28, 2023 •

edited

Loading

DomInvivo commented Aug 28, 2023

DomInvivo Aug 29, 2023

WenkelF Aug 31, 2023 •

edited

Loading

DomInvivo Aug 29, 2023

cwognum Aug 29, 2023

DomInvivo Aug 30, 2023

cwognum Aug 30, 2023

DomInvivo Aug 30, 2023

WenkelF Aug 30, 2023

WenkelF Aug 30, 2023

DomInvivo Aug 31, 2023

WenkelF Aug 31, 2023

WenkelF Aug 31, 2023 •

edited

Loading

DomInvivo Aug 29, 2023

WenkelF Aug 30, 2023

maciej-sypetkowski commented Aug 31, 2023

DomInvivo commented Aug 31, 2023

WenkelF commented Aug 31, 2023

Various updates to graphium #449

Various updates to graphium #449

Conversation

WenkelF commented Aug 28, 2023 • edited Loading

Changelogs

Remarks

codecov bot commented Aug 28, 2023 • edited Loading

Codecov Report

DomInvivo commented Aug 28, 2023

Choose a reason for hiding this comment

WenkelF Aug 31, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WenkelF Aug 31, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

maciej-sypetkowski commented Aug 31, 2023

DomInvivo commented Aug 31, 2023

WenkelF commented Aug 31, 2023

WenkelF commented Aug 28, 2023 •

edited

Loading

codecov bot commented Aug 28, 2023 •

edited

Loading

WenkelF Aug 31, 2023 •

edited

Loading

WenkelF Aug 31, 2023 •

edited

Loading