Exporting nemo nlp models to onnx
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

I am trying to export the nemo_nlp QAModel to onnx format in order to deploy it on triton. I am able to do so using this block of code:
`restored_model = nemo_nlp.models.QAModel.restore_from('finetuning_bert_base_uncased_squad.nemo')
restored_model.export("test_bert_onnx/model.onnx")`

And if it is a larger model (ex. albert-xlarge-v2), I get the protobuf 2GB limit error and to get around this I add _use_external_data_format=True_ to the torch.onnx.export in exportable.py.

This works fine and I get the onnx file and I am able to load it using _onnx.load_ and check it using _onnx.checker.check_model_.
Only after this while running single inference I do not get good results at all. My inference block:
```python
import numpy as np
from onnxruntime import InferenceSession, SessionOptions, get_all_providers
from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def sigmoid(x):
  return 1 / (1 + np.exp(-x))

def create_model_for_provider(model_path: str, provider: str) -> InferenceSession:

  assert provider in get_all_providers(), f"provider {provider} not found, {get_all_providers()}"

  options = SessionOptions()
  options.intra_op_num_threads = 1

  return InferenceSession(model_path, options, providers=[provider])

output_path = "test_bert_onnx/model.onnx"
cpu_model = create_model_for_provider(output_path, "CPUExecutionProvider")

question = "Which NFL team won Super Bowl 50?"
text = "Super Bowl 50 was an American football game to determine the champion of the National Football League (NFL) for the 2015 season. The American Football Conference (AFC) champion Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers 24–10 to earn their third Super Bowl title. The game was played on February 7, 2016, at Levi's Stadium in the San Francisco Bay Area at Santa Clara, California. As this was the 50th Super Bowl, the league emphasized the \"golden anniversary\" with various gold-themed initiatives, as well as temporarily suspending the tradition of naming each Super Bowl game with Roman numerals (under which the game would have been known as \"Super Bowl L\"), so that the logo could prominently feature the Arabic numerals 50."
inputs = tokenizer.encode_plus(question,text, add_special_tokens=True,return_tensors="np")

logits = cpu_model.run(output_names=None,input_feed={"attention_mask":inputs['attention_mask'],"input_ids":inputs["input_ids"],"token_type_ids":inputs["token_type_ids"]})
answer_start_scores = logits[0][0][:,0]
answer_end_scores = logits[0][0][:,1]

ans_tokens = inputs["input_ids"][0][np.argmax(answer_start_scores) : np.argmax(answer_end_scores) + 1 ]
answer_tokens = tokenizer.convert_ids_to_tokens(ans_tokens , skip_special_tokens=True)

answer_tokens_to_string = tokenizer.convert_tokens_to_string(answer_tokens)

print ("\nQuestion : ",question)
print ("\nAnswer : ",answer_tokens_to_string)
```
This works fine but the output answer (and logits) are not what I expect from trained model (The pre-trained bert-base-uncased model that I am using gives an evaluation score of 81.28 exact match and 88.58 f1 score on squad dev1.1 dataset). The output that I get for an example question-context pair taken from squad dataset is this.

![image](https://user-images.githubusercontent.com/72965135/116053402-4652bf80-a698-11eb-9162-03758cb9617d.png)

Can someone point out if I'm missing some processing script of input or output or if the method to export to onnx format is different.
I am able to deploy this in triton but the inference result is same as above.

Lastly, will Triton support NeMo models directly in future as this would prevent a lot of hassle?


**Expected behavior**

Onnx inference working same as _model.inference_ in nemo_nlp

**Environment overview**

 - Environment location: GCP
 - Installed Docker container: docker pull nvcr.io/nvidia/nemo:1.0.0rc1


UPDATE:
I tried all the steps of inference in the [qa_model inference](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/question_answering/qa_model.py#L183) script, here instead of forward pass I did the pass through the onnx inference session and everything worked properly with exact results.
Closing this thread.

I see that this is closed, but I'm having the same issue while trying to serve ONNX on Triton.  Were you able to successfully serve this ONNX model on Triton and get good inference results?  

Yes I was able to serve on Triton and getting inference same as NeMo. @Jblauvs Try switching the attention masks and token type ids inputs. I figured this mismatch in labeling while checking out the ort_inputs sequence in this [tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Text_Classification_Sentiment_Analysis.ipynb).

@avikpal00 Excellent, this worked for me as well.  Thanks for the insight!