ArneBinder / pytorch-ie-hydra-template-1

PyTorch-IE Hydra Template
8 stars 1 forks source link

functionality to use pretrained pie model #168

Closed Bhuvanesh-Verma closed 2 months ago

Bhuvanesh-Verma commented 2 months ago

This PR adds functionality to use a pretrained PyTorch-IE model that allows updating weights of base transformer model. Additionally, it provides an optional parameter, pretrained_pie_model_prefix_mapping, to filter and specify which model layers should be updated. Furthermore, this adds the optional parameter pretrained_pie_model_prefix_mapping that allows to restrict the loading to a subset of parameters via their prefixes.

For instance, to load the weights from a PIE token classification model for further training on conll2003:

python src/train.py experiment=conll2003 pretrained_pie_model_path=path/to/pretrained/pie/model "+pretrained_pie_model_prefix_mapping={model.model:model.model}"

Note that we use model.model:model.model as prefix mapping which will result in only loading the weights from the base transformer model, but not the classification head. This is useful to fine-tune on data with a different set of labels.

In addition, this also adds imports for all models and task modules from the pie_modules package to the train script because these are required when loading a such a model vie pytorch_ie.AutoModel.

Bhuvanesh-Verma commented 2 months ago

Implemented the suggested changes, I think that solved the test failure issue.