Fine-tuning CHGNet to predict materials properties

[yliu1240]

Hi, all!
I'm learning to fine-tune pre-trained CHGNet recently. I'm a little confused how to use the pre-trained CHGNet since I intend to fine-tune CHGNet to predict a few uncommon properties, including several structure/molecular - wise properties (like band gaps) and an atom-wise property (like partial charge, not the original magmoms):
.
.
1) Does the 'state_attr' refer to atom-wise (or to say, site-wise) attribute? From the implementation log #242 , it seems referring to atom-wise attributes, while it seems referring to structure-wise attributes (or to say graph-wise attributes) in tutorial ipynbs (Combining the M3GNet Universal Potential with Property Prediction Models.ipynb, the last section, predicting band gaps). So which one should be the correct interpretation?
.
.
2) To fine-tune pre-trained CHGNets, I guess there must be an atom-wise attribute, right? Because the CHGNet is pre-trained on atom magmoms. Do I need to have an atomic property with a 'magmoms' label as site property in the Structure/Molecule training data? If I intend to fine-tune the pre-trained CHGNet models using other atomic properties, like partial charge instead of 'magmoms', shall I include these site properties in the Structure/Molecule objects in the training data? Can I directly label the site property as 'partial charge' or shall I label it as 'magmoms' (pretending that the partial charges are magmoms). In the fine-tuning process, I guess I must have one atom-wise (site-wise) attribute, but no more than one and no less than one, am I right? Is this because the dimension of atom embeddings are fixed in the pre-training process?
.
.
3) Related to 2), can I fine-tune the pre-trained CHGNets (or any pre-trained MEGNet or M3GNet model) to predict multiple labels (structure/molecule-wise properties other than energies)? It seems I need to remove the final output layer and add another layer for multiple predictions manually by PyTorch? Since the 'ntargets' is fixed in the pre-training process, I guess this can only be achieved by manually loading the pre-trained parameters to another PyTorch NN with output layers with multiple predictions and I didn't find similar APIs in MatGL? Also, when I constructed the MGLDatasets with multiple prediction labels, the fine-tuned models only produced one output.
.
.
4) I'm also struggling to find the suitable collate_fn in my case (predicting multiple structure/molecule-wise targets and an atom-wise target). Collate_fn_graph seems suitable to handle multiple structure-wise labels but can't handle atom-wise predictions. Collate_fn_pes seems designed to handle atom-wise predictions from log #242 and #125 . Do I need to compose a customized collate_fn (probably won't be too difficult as I can copy most from collate_fn_pes)?
.
.
5) Related to 4), in my case, shall I use ModelLightningModule or PotentialLightningModule? It seems only PotentialLightningModule has site-wise (atom-wise) predictions from log #125 (final update message).
.
.
Sorry for long messages and questions here. I'm also trying to find some descriptions and the ipynb (suggested in log #340 )

Thank you so much for your time and patience!

Best regards,

[yliu1240]

It's also interesting to see that, in _m3gnet, _megnet, and _chgnet, the 'state_embedding' is listed with 'node_embedding' and 'edge_embedding', so the 'state' is likely referring to structure-/molecule- wise (graph-wise) features. However, in #242 and in other recent updates, 'state' refers to atom-wise results.

[rydeveraumn]

Is there a notebook that shows how to use CHGNet?