onnx模型转化总结

onnx模型转化总结

最近接手了pytorch模型转onnx的任务，需要将现有pt格式的翻译模型转换为onnx，以方便线上使用。由于机器翻译模型是生成式的，原始方法是将encoder和decoder分别转换为onnx模型，再嵌套beam_search逻辑实现译文生成，但该方法会造成推理速度变慢的问题。因为onnx模型推理是在gpu中的，而beam搜索是在cpu中完成的，因此翻译过程中会反复调用cpu，容易造成延迟问题。由此开发了onnx_for_search脚本，首先对encoder和decoder分别使用jit.trace创建推理图，随后将其嵌入到beam_search逻辑中，使用jit.script对beam函数转换，成功实现了机器翻译的gpu闭环onnx模型实现。
由于onnx不支持torch中的部分方法和操作，因此在onnx模型脚本开发过程中，遇到了很多问题，在此总结以便后续参考。所有问题总结如下：
1、模型主干forward的接受参数应为torch.tensor，forward中如需调用其他函数，实参传入时也应转换为tensor，即便实参只是单个int或float变量；
2、函数调用不能传入encoder或decoder模型，按照1的说法，传入格式必须为tensor，模型传入后也会被当作tensor处理，发生意外错误；
3、不支持字典、列表等格式，不能在函数中定义字典或列表，但调用另一个函数时，该函数内可以；
4、torch.repeat、torch.new_tensor、torch.new_full方法均不支持，需要使用cat、tensor、full等操作替换；
5、具备size属性的空张量，无法在函数中做reshape操作；
6、子调用函数中如果有list，不支持对其中元素赋值，但可以拼接或分片。