Finetuning pipeline

[WenkelF]

@DomInvivo as discussed a first draft for the Finetuning pipeline.

Two possible pipelines:
expts/main_run_finetuning_v1.py (probably to be removed):

• Start with pretrained model
• Remove redundant task heads
• Drop/add layers of model part to be finetuned
• Drawback: Difficult updating of hyperparameters

expts/main_run_finetuning_v2.py:

• Initialize new model with final finetuning architecture
• Overwrite parameters that are shared with pretrained model

Main TODOs:

• Use hydra to easy finetuning config
• Create unit test with dummy pretrained model

[codecov]

Codecov Report

Merging #414 (6625e0c) into main (9515ad6) will decrease coverage by 2.14%.
The diff coverage is 20.51%.

@@            Coverage Diff             @@
##             main     #414      +/-   ##
==========================================
- Coverage   66.87%   64.74%   -2.14%     
==========================================
  Files          82       89       +7     
  Lines        7838     8211     +373     
==========================================
+ Hits         5242     5316      +74     
- Misses       2596     2895     +299     

Flag	Coverage Δ
unittests	64.74% <20.51%> (-2.14%)	⬇️

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

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

[DomInvivo]

Here's my review. It mostly aligns with what we discussed last time

[DomInvivo]

This could be a function within FeedForwardNN and FullGraphNetwork

[WenkelF]

I am not sure, we can remove the function and replace inside the networks. As discussed in that #411, it might be possible to have a similar function within load_architecture in graphium/data/_loader.py but depends on the PR.

[cwognum]

Hi @WenkelF, sorry for the delay here!

Thank you for this first implementation!

I left comments whenever something stood out to me, but I'm aware that this PR is still WIP. Sorry if I pointed out some things that you were already planning to change.

It would be super helpful if you could document the main fine-tuning "flow". Furthermore, I would suggest to simplify the process by adding support for one feature at a time, instead of having half-implemented features. This will make understanding, debugging, testing and maintaining the code base a lot easier.

Happy to help next week!

[DomInvivo]

Don't forget that the objective is to release a working, incomplete version first. Then refining it with more complex fine-tuning possibilities.

[DomInvivo]

This doesn't make sense. We should not have any architectural choice from the original pre-trained model in here. Only things that would change.

That way, we can take different pre-trained models that have different hparams/seed and fine-tune them all with the same file.

[WenkelF]

I agree. The configurations are still structured in a way where we have access to both the full config of the pretrained model and the pretraining-related config. And the modify_cfg_for_finetuning function consolidates information to one config.

This will be fixed once we incorporate the new hydra config from #421. We will still need modify_cfg_for_finetuning as of now. Therefore, it could be good waiting for the final version.

[DomInvivo]

I don't even understand the point of having a try and an if there. Dict.pop works even if the key is not available.

But we need to make sure that args[task] is not used elsewhere, even outside the current function since dict are passed as pointers. We only want to remove epoch_sampling for the hash key. So I would suggest the following.

Suggested change

	try:
	if "epoch_sampling_fraction" in args[task].keys():
	args[task].pop("epoch_sampling_fraction")
	except:
	pass
	args[task] = deepcopy(args[task])
	args[task].pop("epoch_sampling_fraction")

[WenkelF]

Indeed, this is not an ideal fix. I wanted to investigate the issue a bit more.

We cannot use args[task].pop("epoch_sampling_fraction") because args[task] may be of class DatasetProcessingParams instead of Dict. In particular, ADMETBenchmarkDataModule makes use of DatasetProcessingParams.

The error originates from changes here 4b82ba3, where the line args[task].pop("epoch_sampling_fraction") was added. It did not cause errors back then because we were using a Dict in all configs (although I see a comment # To be replaced by a new class "DatasetParams" everywhere it appears).

Will create issue and think about a fix.

[DomInvivo]

Change to explicitly pass parameters.

Suggested change

	def __init__(self, cfg, train_bn: bool = False):
	def __init__(self, fine-tuning, architecture, module_from_pretrained, ....................., train_bn: bool = False):

[WenkelF]

Addressed in 76e2ba6

[DomInvivo]

else: raise "Wrong module"

[WenkelF]

Addressed in 76e2ba6

[DomInvivo]

I would define all these in a dictionary _module_map = {pe_encoders: pl_module.model.encoder_manager, ...} directly in the __init__. That way, with inheritance, someone could modified the entries without copy-pasting all the logic.

_module_map can replace the module_list you already have.

But instead of a regular dict, using an OrderedDict would also allow you to say something like: "freeze everything before task_heads" in a very simple way.

[DomInvivo]

It's a bad idea to have a FullGraphFinetuningNetwork that basically copy-pastes most of the functionality of FullGraphNetwork.

Either use inheritance, or implement the fine-tuning logic directly within FullGraphNetwork

[WenkelF]

@DomInvivo thanks for your comments. Here is also a quick overview of the updates

Updates:

• Moved finetuning-related architecture out of FullGraphMultitaskNetwork to FullGraphFinetuningNetwork under graphium/finetuning/finetuning_architecture
  • More flexibility to build finetuning head on top of pretrained model
  • For now, finetuning is specific to FullGraphMultitaskNetwork (as pretrained model), but we can finetune from gnn, graph_output_nn and task_heads now
  • Additionally, we can add a flexible FinetuningHead, inheriting from nn.Module and MupMixin. For example, it can be a FeedForwardNN, FeedForwardPyg, TaskHeads, or a custom network that can be added to FINETUNING_HEAD_DICT
  • FinetuningHead needs a bit more work to correctly support all possible scenarios (modules/levels to finetune from); if not used, we automatically fall back to the finetuning logic, where we can only drop & add depth to the module we finetune from
• Finished functions for overwriting shared weights and partially freezing modules while training and extended them from only task_heads to also gnn and graph_output_nn
• Added preliminary unit test the discussed “minimal” finetuning pipeline (finetune from one of the task heads)
  • We currently test depths, in_dim, ou_dim of changed modules are correct and overwritten correctly
  • Test does not cover training/freezing yet
  • Planning to add tests for finetuning from other modules as well
• Configuration has been integrated with hydra but needs more work to be easier to use (@luis-mueller is helping with that)

Remarks:

• Hidden_dims can only be used as int together with depth at the moment
• Currently, everything has only been tested on cpu/gpu

[WenkelF]

@DomInvivo this pull introduces the following:

• Finetuning pipeline
• Updates to hydra configuration (by @cwognum )

The finetuning pipeline is maintained separately from existing architectures under graphium/finetuning/finetuning_architecture. The FullGraphFinetuningNetwork comes with two submodules:

• PretrainedModel
• FinetuningHead (optional)

PretrainedModel can load pretrained models (e.g., FullGraphMultitaskNetwork) and allows basic finetuning operations (in particular from the task heads), including dropping or extending layers of the module to finetune from.

FinetuningHead is optional and allows to define fully customizable networks that are applied on top of PretrainedModel. If not specified, we fall back to basic finetuning explained above.

Training is handled by the GraphFinetuning callback in graphium/finetuning/finetuning.

All methods in graphium/finetuning are implemented such that they are not specific to a pretrained model or finetuning head. This is achieved by requiring the pretrained model to come with a module_map (see, e.g., FullGraphMultitaskNetwork) that facilitates setting pretrained weights and freezing the correct layers.

The new unit test tests/test_finetuning tests the pipeline end-to-end for a specific case. Together with the corresponding config, it might be a good starting point to understand the pipeling and what information is needed in the config.

The updates to hydra allow to easily switch between benchmarking and finetuning. Major changes are documented here WenkelF#10

[DomInvivo]

Mostly looks good, great work on this major PR!
A few changes to make, and some comments.

[DomInvivo]

I don't think this should be here. I think that, if you are using a pre-trained model, you should pass directly model_kwargs["pretrained_model_kwargs"] into model_kwargs.

[WenkelF]

Removed in a3d4715 as explained below

[DomInvivo]

In general, the PredictorModule should be agnostic to the model passed. By having the self._get_task_key here, it forces a certain architecture in the config which is not very flexible.

I see that this logic was introduced already in the code prior to this PR. If it requires too many changes, let's open a new issue.

[WenkelF]

You are right, thanks for pointing out.

We could achieve this by getting the task-specific information (which is only the task level as far as I know) from the datamodule.

Here is how this could be done:
a3d4715

What do you think?

[DomInvivo]

Again, don't like how the config structure is imposed. Perhaps task_level should simply be passed to the PredictorModule to keep flexibility

[WenkelF]

Good idea, this is implemented in a3d4715

[DomInvivo]

Why duplicating the model for CPU and GPU? Models should be agnostic to the training hardware, and to the fine-tuning hardware. You can train on CPU and fine-tune on GPU or IPU

[WenkelF]

Yes, I agree. I only changed from gpu to cpu because github cannot do unit tests on gpu. Should I remove the gpu model?

[DomInvivo]

Yes, you can remove the gpu model.

[DomInvivo]

@zhiyil-graphcore @s-maddrellmander We'll need your help here to fix the tests for IPU. And ideally, have a test that loads a CPU-trained model onto IPU for finetuning.

[WenkelF]

@cwognum the bug in the finetuning training is fixed here febdf2d

I missed a deepcopy operation when defining the datahashes for the TDC datasets. We include the first 5 rows of the df when generating the hash and the bug reduced the datasets to those 5 rows (molecules) as well.

Make sure to remove the TDC datasets from datacache. You can use datamodule._path_to_load_from_file("train) to get "train_[data_hash]" and then remove all data hashes (e.g., "train_[data_hash]", "test_[data_hash]", etc. that end with the same sequence.

[WenkelF]

@DomInvivo I added some final improvements 24354ee

1. Made modify_cfg_for_finetuning function fully agnostic to pretrained model
2. Dropped dummy-pretrained-model-gpu and added map_location as an argument to load_pretrained_models function
3. Added the option to keep modules after finetuning module

(3.) makes it much easier to finetune from modules other than the task_heads without manually re-defining the downstream network. When finetuning from task_heads, it is not needed.

[s-maddrellmander]

@WenkelF - try merging from the main branch, I've made a change to the IPU test CLI that should take into account the changes made in this PR. If that doesn't work let me know.