Pretraining

[IHumonen]

Hello!
I have a model (allegro), trained via nequip, and now I want to retrain this model on another dataset. Is it possible to do such a thing using nequip? Thanks in advance for the answer.

[Linux-cpp-lisp]

Hi @IHumonen ,

This is definitely possible, and is one of the main uses of the initialize_from_state model builder (see docstring

nequip/nequip/model/_weight_init.py

Line 13 in b62c4f4

def initialize_from_state(config: Config, model: GraphModuleMixin, initialize: bool):

).

Basically:

1. Copy your config file and give it a new run name for the new training (the finetuning)
2. Change the dataset and any other trainer settings that you want to update for your finetuning/further training
3. Add the initialize_from_state model builder to the end of your model_builders list
4. Add the initial_model_state key to your config and set it to the path to a state dict of saved weights from which to initialize the model, such as the best_model.pth weights from your previous training run:

   initial_model_state: /path/to/pretraining/session/best_model.pth
   

[Linux-cpp-lisp]

@wwwjjj123yyyy can you please post a complete stack trace?

[wwwjjj123yyyy]

This is the bottom lines of my output file when running nequip, I hope it helps.

[Linux-cpp-lisp]

I can't open the file; can you please provide it as text here in GitHub?

[wwwjjj123yyyy]

@Linux-cpp-lisp This is my otuput file.
slurm-286959.txt

[Mathieu-Istas]

Hi, I do have the same error as @wwwjjj123yyyy . I copied a config I used to train a model and wanted to do some training with some extra data with it, so I added the lines

### start from a previous run
model_builders:
 - initialize_from_state
initial_model_state: path_to_previous_model/best_model.pth

And after processing the data I got the error

Traceback (most recent call last):
  File "/home/istas/.local/bin/nequip-train", line 8, in <module>
    sys.exit(main())
  File "/home/istas/.local/lib/python3.9/site-packages/nequip/scripts/train.py", line 83, in main
    trainer = fresh_start(config)
  File "/home/istas/.local/lib/python3.9/site-packages/nequip/scripts/train.py", line 211, in fresh_start
    final_model = model_from_config(
  File "/home/istas/.local/lib/python3.9/site-packages/nequip/model/_build.py", line 141, in model_from_config
    model = GraphModel(
  File "/home/istas/.local/lib/python3.9/site-packages/nequip/nn/_graph_model.py", line 47, in __init__
    self._init_irreps(irreps_in=irreps_in, irreps_out=model.irreps_out)
AttributeError: 'NoneType' object has no attribute 'irreps_out'

Should I specify irreps_out manually in the config file? Something like chemical_embedding_irreps_out: 32x0e?

[IZugec]

@Linux-cpp-lisp
Firstly, thanks for clarifications and the discussion. I wanted to talk about the statement you have written above:

"Be careful not to change any model hyperparameters in the finetuning YAML that would affect the meaning of the weights."

I wanted to ask you specifically for the radius which is obviously extremely important (here I am retraining nequip). Do you think it's safe to change radius as a hyperparameter?

[Linux-cpp-lisp]

Hi @IZugec ,

Let me make sure I understand what you're asking before I answer: you want to train a model with some r_max, and then fine-tune it later with a different r_max? Can you explain a little bit the motivation, if this is the case, because it doesn't make a ton of sense to me on its own.

[IZugec]

You understood correctly. I want to fine-tune it on different dataset which I think would benefit from larger r_max, but want to continue from weights that was learned on this initial dataset.

[Linux-cpp-lisp]

I see... this is complicated by the radial basis projection: our networks don't process distances directly, but rather project them out into a radial basis (we use the Bessel basis, see the papers for details and references). This radial basis depends on r_max and the number of radial basis functions, so if you change r_max the basis changes and thus completely changes the interpretation of the weights.

The first approach that comes to my mind would be to set a larger r_max for the radial basis in the original model using:

BesselBasis_r_max: 10.0  # or whatever

but still setting a smaller overall r_max. This way the original dataset and smooth cutoff will be set to your smaller cutoff. Then, when you want to fine tune, you also set r_max to 10.0 (or whatever), but the radial basis won't change--- only the polynomial cutoff and the dataset preprocessing cutoff, so the fine-tuning should go OK...

[IZugec]

Hmm I see. Thank you very much for your time to reply. Appreciate the insight.

[ken1944]

Any suggestions to solve CUDA out of memory error while finetuning? I am using a batch size of 1.

[Linux-cpp-lisp]

Is there a specific reason you think this is related to finetuning vs training in general? Can you please open a new discussion thread with a full description and traceback?

[ken1944]

Sure, made a new one here #388