How to extract features from a pytorch pretrained fine-tuned model

Question

I need to extract features from a pretrained (fine-tuned) BERT model.

I fine-tuned a pretrained BERT model in Pytorch using huggingface transformer. All the training/validation is done on a GPU in cloud.

At the end of the training, I save the model and tokenizer like below:

best_model.save_pretrained('./saved_model/')
tokenizer.save_pretrained('./saved_model/')

This creates below files in the saved_model directory:

config.json
added_token.json
special_tokens_map.json
tokenizer_config.json
vocab.txt
pytorch_model.bin

I save the saved_model directory in my computer and load the model and tokenizer like below

model = torch.load('./saved_model/pytorch_model.bin',map_location=torch.device('cpu'))
tokenizer = BertTokenizer.from_pretrained('./saved_model/')

Now to extract features, I do below

input_ids = torch.tensor([tokenizer.encode("Here is some text to encode", add_special_tokens=True)])
last_hidden_states = model(input_ids)[0][0]

But for the last line, it throws me error TypeError: 'collections.OrderedDict' object is not callable

It seems like I am not loading the model properly. Instead of loading the entire model in itself, I think my model=torch.load(....) line is loading a ordered dictionary.

What am I missing here? Am I even saving the model in the right way? Please suggest.

Answer 1

torch.load() returns a collections.OrderedDict object. Checkout the recommended way of saving and loading a model's state dict.

Save:

torch.save(model.state_dict(), PATH)

Load:

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

So, in your case, it should be:

model = BertModel(config)
model.load_state_dict('./saved_model/pytorch_model.bin',
                       map_location=torch.device('cpu'))
model.eval() # to disable dropouts