如何从检查点文件加载经过微调的 pytorch huggingface bert model？

Question

I had fine tuned a bert model in pytorch and saved its checkpoints via torch.save(model.state_dict(), 'model.pt')

Now When I want to reload the model, I have to explain whole network again and reload the weights and then push to the device.

Can anyone tell me how can I save the bert model directly and load directly to use in production/deployment?

Following is the training code and you can try running there in colab itself, After training completion. you will notice in file system we have a checkpoint file. But I want to save the model itself.

LINK TO COLAB NOTEBOOK FOR SAMPLE TRAINING

Following is the current inferencing code I written.

import torch
import torch.nn as nn
from transformers import AutoModel, BertTokenizerFast
import numpy as np
import json

tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
device = torch.device("cpu")


class BERT_Arch(nn.Module):

    def __init__(self, bert):
        super(BERT_Arch, self).__init__()

        self.bert = bert

        # dropout layer
        self.dropout = nn.Dropout(0.1)

        # relu activation function
        self.relu = nn.ReLU()

        # dense layer 1
        self.fc1 = nn.Linear(768, 512)

        # dense layer 2 (Output layer)
        self.fc2 = nn.Linear(512, 2)

        # softmax activation function
        self.softmax = nn.LogSoftmax(dim=1)

    # define the forward pass
    def forward(self, sent_id, mask):
        # pass the inputs to the model
        _, cls_hs = self.bert(sent_id, attention_mask=mask, return_dict=False)

        x = self.fc1(cls_hs)

        x = self.relu(x)

        x = self.dropout(x)

        # output layer
        x = self.fc2(x)

        # apply softmax activation
        x = self.softmax(x)

        return x


bert = AutoModel.from_pretrained('bert-base-uncased')
model = BERT_Arch(bert)
path = './models/saved_weights_new_data.pt'
model.load_state_dict(torch.load(path, map_location=device))
model.to(device)

def inference(comment):
    tokens_test = tokenizer.batch_encode_plus(
        list([comment]),
        max_length=75,
        pad_to_max_length=True,
        truncation=True,
        return_token_type_ids=False
    )
    test_seq = torch.tensor(tokens_test['input_ids'])
    test_mask = torch.tensor(tokens_test['attention_mask'])
    predictions = model(test_seq.to(device), test_mask.to(device))
    predictions = predictions.detach().cpu().numpy()
    predictions = np.argmax(predictions, axis=1)
    return predictions

I simply want to save a model from this notebook in a way such that I can use it for inferencing anywhere.

Answer 1

Just save your model using model.save_pretrained, here is an example:

model.save_pretrained("<path_to_dummy_folder>")

You can download the model from colab, save it on your gdrive or at any other location of your choice. While doing inference, you can just give path to this model (you may have to upload it) and start with inference.

To load the model

model = AutoModel.from_pretrained("<path_to_saved_pretrained_model>")

#Note: Instead of AutoModel class, you may use the task specific class as well.