Adding linear layer to transformer model (+ save_pretrained and load_pretrained)
I want to extend a transformer model (let’s take bert, electra, etc for example) with a linear layer and initialize the linear layer with the same initializer as the transformer model. I also want save_pretrained and load_pretrained to work smoothly (i.e. save and load the model WITH the linear layer, not separately). This is what I’m currently doing, but it’s not workingclass ExtendedTransformer(PreTrainedModel):
    base_model_prefix = "model"
    def __init__(self, config):
        super().__init__(config)
        self.model = AutoModel.from_pretrained('bert-base-uncased')
        self.linear = torch.nn.Linear(self.config.hidden_size, 128, bias=False)
        self.model.init_weights()Any advice on how to do this properly?Thanks
2022-03-10T08:37:24Z
If you just want to increase the output dimensions, you can simply usemodel = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=128)But here’s an explanation of what I think the issue is with your code.If I’m not mistaken, your code it only initializes themodel(which is already being initialized by thefrom_pretrainedcall). As far as I know, you should sub-class specific architectures, becausePreTrainedModelitself does not have aninit_weightsmethod and models may implement a different scheme for that. Since the end of last year, it’s probably best to usepost_initinstead ofinit_weights, too. Maybe you can achieve this for a variety of models with AutoModel as well, but I am not sure.So, the following should work. It is basically the same asBertForSequenceClassification.from transformers import BertPreTrainedModel, AutoConfig, AutoModel
import torch

class ExtendedTransformer(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.bert = AutoModel.from_pretrained('bert-base-uncased')
        self.linear = torch.nn.Linear(self.config.hidden_size, 128, bias=False)
        self.post_init()

if __name__ == '__main__':
    config = AutoConfig.from_pretrained('bert-base-uncased')
    inst = ExtendedTransformer(config)