How to Export Stanza to ONNX format?

Question

How to export Stanza to ONNX format? It seems impossible to just simply train the model.

Answer 1

There is an explanation here: https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

I created a fork from stanza for this experiment here https://github.com/vivkvv/stanza. See also my commits https://github.com/vivkvv/stanza/commits?author=vivkvv.

I used pipeline_demo.py for testing. The main thing I added is code just inside models/tokanization/trainer.py below the line 77

pred = self.model(units, features)

Due to explanation I added

        torch.onnx.export(
            self.model,
            (units, features),
            onnx_export_file_name,
            opset_version=9,
            export_params=True,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        )

and it works for tokenization. But the same does not work for e.g. pos or lemmatizer (see my commit for PartOfSpeech). And I get different errors for different opset_version.

I created a question on github/stanza and you could see there https://github.com/stanfordnlp/stanza/issues/893