Converting biomegatron to onnx

Open ayush-raj7 opened this issue 3 years ago • 0 comments

Describe the bug

Hi, I am trying to convert Biomegatron Model finetuned for token classification tasks using nemo to onnx format. I am getting the following error:

ARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. WARNING: The shape inference of prim::PythonOp type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. [NeMo E 2022-08-08 14:55:29 export:164] Export failed. Please make sure your NeMo model class (<class 'nemo.collections.nlp.models.token_classification.token_classification_model.TokenClassificationModel'>) has working export() and that you have the latest NeMo package installed with [all] dependencies. Traceback (most recent call last): File "export.py", line 173, in nemo_export(sys.argv[1:]) File "export.py", line 169, in nemo_export raise e File "export.py", line 155, in nemo_export _, descriptions = model.export( File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/exportable.py", line 68, in export out, descr, out_example = model._export( File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/exportable.py", line 164, in _export torch.onnx.export( File "/opt/conda/lib/python3.8/site-packages/torch/onnx/init.py", line 332, in export return utils.export(model, args, f, export_params, verbose, training, File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 125, in export _export(model, args, f, export_params, verbose, training, input_names, output_names, File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 824, in _export proto, export_map, val_use_external_data_format, node_names = graph._export_onnx( RuntimeError: ONNX export failed: Couldn't export Python operator ScaledMaskedSoftmax

Steps/Code to reproduce bug I'm using the script export.py used in the below tutorial: [(https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/MegatronBert_export.ipynb)]

Please, download export.py from here: [(https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/export.py)]

Also, tried model.export('model.onnx') separately and getting the same error.

Environment overview (please complete the following information)

Environment location: Docker, GCP
docker pull nvcr.io/nvidia/nemo:22.05

Environment details

NeMo container: nvcr.io/nvidia/nemo:22.05 NeMo version: 1.10.0 Branch: 'Main' Biomegatron Model: "biomegatron345m_biovocab_50k_cased"

Aug 08 '22 15:08 ayush-raj7