如何将Stanza模型导出为ONNX格式？

Stanza底层基于PyTorch开发，直接导出整个Pipeline失败的核心原因是模型中存在大量动态控制流、Python原生实现的后处理逻辑，以及未适配ONNX规范的自定义算子，按以下步骤操作即可完成导出：

• 第一步：拆解Pipeline单独导出组件
  不要直接导出完整的Stanza Pipeline对象，Stanza的Pipeline是多个模型组件+Python后处理逻辑的封装，无法直接转为ONNX计算图。你需要单独取出分词、词性标注、依存句法分析等单个组件的底层PyTorch模型分别导出。
• 第二步：适配模型逻辑满足ONNX导出要求
  关闭Stanza模型中所有动态分支逻辑：固定推理时的最大序列长度、关闭动态padding开关、注释掉根据输入维度自适应调整计算逻辑的if-else分支，避免ONNX导出时无法解析动态控制流。
• 第三步：执行导出操作并配置动态轴
  NLP场景下输入的batch大小、序列长度都是可变的，导出时必须显式指定动态轴，否则导出的模型仅支持固定维度输入，参考代码如下：

  import stanza
  import torch
  
  # 以导出中文分词模型为例
  nlp = stanza.Pipeline(lang='zh', processors='tokenize')
  tokenizer_model = nlp.processors['tokenize'].trainer.model
  # 必须切换到推理模式再导出
  tokenizer_model.eval()
  
  # 构造和实际推理维度对齐的样例输入
  sample_input = torch.randint(0, 21128, (1, 32), dtype=torch.long)
  
  torch.onnx.export(
      model = tokenizer_model,
      args = sample_input,
      f = "stanza_tokenizer.onnx",
      export_params = True,
      opset_version = 17, # 建议用17及以上版本，支持更多NLP相关算子
      do_constant_folding = True,
      dynamic_axes = {
          "input": {0: "batch_size", 1: "seq_len"},
          "output": {0: "batch_size", 1: "seq_len"}
      },
      input_names = ["input"],
      output_names = ["output"]
  )
  

• 第四步：验证导出模型正确性
  用ONNXRuntime加载导出的模型，和原PyTorch模型的输出做误差校验，误差小于1e-5即可认为导出正常，参考代码如下：

  import onnxruntime as ort
  import numpy as np
  
  ort_session = ort.InferenceSession("stanza_tokenizer.onnx")
  ort_output = ort_session.run(None, {"input": sample_input.numpy()})[0]
  torch_output = tokenizer_model(sample_input).detach().numpy()
  
  assert np.allclose(torch_output, ort_output, atol=1e-5), "模型输出误差超出阈值，导出失败"
  

注意事项 #

• 若需要完整的端到端推理能力，需将多个导出的ONNX模型和对应的后处理逻辑用ONNXRuntime/Numpy重新实现，Stanza原生的Python后处理逻辑无法被导出到ONNX计算图中。
• 部分小语种模型用到了自定义PyTorch算子，导出前需要将自定义算子转换为ONNX支持的原生算子，或自行注册ONNX自定义算子，否则会触发导出报错。
• 不要在模型训练状态下直接导出，必须加载训练好的权重并切换到eval模式后再执行导出操作，否则训练模式下的Dropout、BatchNorm层会导致导出模型的输出不符合预期。

内容的提问来源于stack exchange，提问作者George Clooney