Clarification on --pretrained usage (for fine-tuning?) and possible parser issue

[eleproia]

Dear all,

First of all, congratulations on this repository — great work!

We would like to fine-tune GEMS on proprietary crystal structures, starting from an already trained GEMS checkpoint and using it as a pretrained model.

While looking into train.py, we had a couple of questions regarding the --pretrained argument:

1. Is using the --pretrained flag to load a model checkpoint and train the model with additional crystal structures the intended and correct way to use this option?

2. We noticed that the argument parser currently defines

parser.add_argument(
    "--pretrained",
    default=False,
    type=lambda x: x.lower() in ['true', '1', 'yes'],  
    help="Provide the path of a state dict that should be imported"
)

However, this seems to cast the argument to a boolean, while the help string (and the code below) suggests that a string path to a checkpoint should be provided.
Would it be more appropriate for --pretrained to be of type str, so that it can directly accept the path to the model checkpoint?

Thank you very much for your time and for maintaining this project.

Best regards,
Eleonora

[DavidGraber]

Dear Eleonora

Sorry for the delay in responding, I was out-of-office for the first months of the year. I hope that GEMS was useful for you in the meantime, and I thank you for pointing out this clear error in my script. I will have a look at how to improve it now. Actually I have never really used this parser argument and just added this functionality for the sake of completeness.

I have never tried your use case, but given your dataset of crystal structures is sufficiently large, I would try it like you suggested and maybe play around with the learning rate a bit.

Best wishes,
David