使用ONNX搭建NLP Transformers pipelines

大数据文摘授权转载自数据派THU

$ conda create -y -n hf-onnx python=3.8$ conda activate hf-onnx$ git clone https://github.com/ChainYo/transformers-pipeline-onnx.git$ cd transformers-pipeline-onnx$ pip install -r requirements.txt

from transformers.onnx.features import FeaturesManagerdistilbert_features = list(FeaturesManager.get_supported_features_for_model_type("bert").keys())print(distilbert_features)>>> ['default', 'masked-lm', 'causal-lm', 'sequence-classification', 'token-classification', 'question-answering']

$ python -m transformers.onnx \    --model=dslim/bert-base-NER \    --feature=token-classification \    onnx/

import torchfrom onnxruntime import (InferenceSession, SessionOptions, GraphOptimizationLevel)from transformers import (TokenClassificationPipeline, AutoTokenizer, AutoModelForTokenClassification)

options = SessionOptions() # initialize session optionsoptions.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALLsession = InferenceSession("onnx/model.onnx", sess_options=options, providers=["CPUExecutionProvider"])# disable session.run() fallback mechanism, it prevents for a reset of the execution providersession.disable_fallback()

from onnxruntime import get_all_providersget_all_providers()

class OnnxTokenClassificationPipeline(TokenClassificationPipeline):def __init__(self, *args, **kwargs):super().__init__(*args, **kwargs)def _forward(self, model_inputs):"""Forward pass through the model. This method is not to be called by the user directly and is only usedby the pipeline to perform the actual predictions.This is where we will define the actual process to do inference with the ONNX model and the session createdbefore."""# This comes from the original implementation of the pipelinespecial_tokens_mask = model_inputs.pop("special_tokens_mask")offset_mapping = model_inputs.pop("offset_mapping", None)sentence = model_inputs.pop("sentence")inputs = {k: v.cpu().detach().numpy() for k, v in model_inputs.items()} # dict of numpy arraysoutputs_name = session.get_outputs()[0].name # get the name of the output tensorlogits = session.run(output_names=[outputs_name], input_feed=inputs)[0] # run the sessionlogits = torch.tensor(logits) # convert to torch tensor to be compatible with the original implementationreturn {"logits": logits,"special_tokens_mask": special_tokens_mask,"offset_mapping": offset_mapping,"sentence": sentence,**model_inputs,}# We need to override the preprocess method because the onnx model is waiting for the attention masks as inputs# along with the embeddings.def preprocess(self, sentence, offset_mapping=None):truncation = True if self.tokenizer.model_max_length and self.tokenizer.model_max_length > 0 else Falsemodel_inputs = self.tokenizer(sentence,return_attention_mask=True, # This is the only difference from the original implementationreturn_tensors=self.framework,truncation=truncation,return_special_tokens_mask=True,return_offsets_mapping=self.tokenizer.is_fast,)if offset_mapping:model_inputs["offset_mapping"] = offset_mappingmodel_inputs["sentence"] = sentencereturn model_inputs

model_name_from_hub = "dslim/bert-base-NER"tokenizer = AutoTokenizer.from_pretrained(model_name_from_hub)model = AutoModelForTokenClassification.from_pretrained(model_name_from_hub)onnx_pipeline = OnnxTokenClassificationPipeline(task="ner",model=model,tokenizer=tokenizer,framework="pt",aggregation_strategy="simple",)

sequence = "Apple was founded in 1976 by Steve Jobs, Steve Wozniak and Ronald Wayne to develop and sell Wozniak's Apple I personal computer"onnx_pipeline(sequence)

pytorch_pipeline = TokenClassificationPipeline(task="ner",model=model,tokenizer=tokenizer,framework="pt",aggregation_strategy="simple",)

sequences = {"short_sequence": "Hello my name is Thomas and I love HuggingFace.","medium_sequence": "Winston Churchill was born in 1874 in Stoke-on-Trent, England, to a German father, William and Elizabeth Churchill.","long_sequence": """The first person to reach the summit of Everest was the South Nepalese Everest Gurun, who was a member of the Royal Nepal Expedition, led by the Nepalese Mountaineer, Sir Edmund Hillary. Hilary lived in the Himalayas for a time. He sadly died in 1953 at the age of 88."""}

import timeitfrom tabulate import tabulateresults = [["Sequence Length", "PyTorch", "ONNX"]]for k, v in sequences.items():results.append( [k, timeit.timeit(lambda: pytorch_pipeline(v), number=300), timeit.timeit(lambda: onnx_pipeline(v), number=300)])print(tabulate(results, headers="firstrow")

print(f"For a short sequence: ONNX is {results[1][1]/results[1][2]:.2f}x faster than PyTorch")print(f"For a medium sequence: ONNX is {results[2][1]/results[2][2]:.2f}x faster than PyTorch")print(f"For a long sequence: ONNX is {results[3][1]/results[3][2]:.2f}x faster than PyTorch")