如何将预训练KoGPT2模型导出为ONNX格式，解决转换报错问题？

1. 核心报错由PyTorch 1.6.0版本的已知ONNX导出缺陷导致，Unable to cast from non-held to held instance是该版本处理带动态索引的计算图时的特有bug，GPT2模型的因果掩码计算逻辑刚好触发了该问题。
2. 依赖冲突也会加剧导出异常：你同时安装了正式版onnxruntime和测试版ort-nightly，二者存在功能冲突；同时旧版本transformers的convert_graph_to_onnx工具对GPT类模型的适配存在缺失。
3. 转换代码未指定pipeline类型，默认加载的基础GPT2Model会丢弃lm_head权重，不符合后续转TensorFlow做推理的需求。

第一步：清理并安装兼容依赖 #

!pip uninstall -y torch torchvision onnxruntime ort-nightly onnxruntime-tools
!pip install torch==1.9.0+cpu torchvision==0.10.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
!pip install onnxruntime==1.11.0 transformers==4.20.1 onnx==1.12.0

以上组合为经过验证的KoGPT2转ONNX的兼容依赖组合。

第二步：修改转换代码 #

!rm -rf onnx/
from pathlib import Path
from transformers.convert_graph_to_onnx import convert

convert(
    framework="pt", 
    model="skt/kogpt2-base-v2", 
    output=Path('/content/drive/MyDrive/kogptonnx/kogpt.onnx'), 
    opset=12,
    pipeline_name="text-generation"
)

指定pipeline_name="text-generation"会自动加载带lm_head的完整GPT2LMHeadModel，保留文本生成所需的输出头，避免权重丢失问题。

按照上述方案操作即可正常导出KoGPT2的ONNX格式模型，导出后的模型可直接通过onnxruntime推理，也可通过转换工具转为TensorFlow格式使用。

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