Converting a PyTorch BERT-NER Model¶
The goal of this article is to present a step-by-step guide on how to convert PyTorch BERT-NER model to OpenVINO IR. First, you need to download the model and convert it to ONNX.
Downloading and Converting the Model to ONNX¶
To download a pretrained model or train the model yourself, refer to the instructions in the BERT-NER model repository. The model with configuration files is stored in the out_base
directory.
To convert the model to ONNX format, create and run the following script in the root directory of the model repository. If you download the pretrained model, you need to download bert.py to run the script. The instructions were tested with the commit-SHA: e5be564156f194f1becb0d82aeaf6e762d9eb9ed
.
import torch
from bert import Ner
ner = Ner("out_base")
input_ids, input_mask, segment_ids, valid_positions = ner.preprocess('Steve went to Paris')
input_ids = torch.tensor([input_ids], dtype=torch.long, device=ner.device)
input_mask = torch.tensor([input_mask], dtype=torch.long, device=ner.device)
segment_ids = torch.tensor([segment_ids], dtype=torch.long, device=ner.device)
valid_ids = torch.tensor([valid_positions], dtype=torch.long, device=ner.device)
ner_model, tknizr, model_config = ner.load_model("out_base")
with torch.no_grad():
logits = ner_model(input_ids, segment_ids, input_mask, valid_ids)
torch.onnx.export(ner_model,
(input_ids, segment_ids, input_mask, valid_ids),
"bert-ner.onnx",
input_names=['input_ids', 'segment_ids', 'input_mask', 'valid_ids'],
output_names=['output'],
dynamic_axes={
"input_ids": {0: "batch_size"},
"segment_ids": {0: "batch_size"},
"input_mask": {0: "batch_size"},
"valid_ids": {0: "batch_size"},
"output": {0: "output"}
},
opset_version=11,
)
The script generates ONNX model file bert-ner.onnx
.
Converting an ONNX BERT-NER model to IR¶
mo --input_model bert-ner.onnx --input "input_mask[1,128],segment_ids[1,128],input_ids[1,128]"
where 1
is batch_size
and 128
is sequence_length
.