代码中使用了bitpack来量化-反量化,其中用到了torch.bitwise_right_shift和bitwise_and函数,在导出ONNX模型时报错,特此记录。
1. Opset version选择
从ONNX DOCUMENT可以看到,onnx直到opset 18才添加了对BitwiseAnd的支持。
2. 直接导出,报错
导出脚本:
torch.onnx.export(model, (x), "test.onnx", opset_version=18, verbose=True, input_names=['input'])
运行后:
raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::bitwise_right_shift' to ONNX opset version 18 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.
可以看到虽然ONNX 官方说了支持,但并没有与torch函数对应起来。
3. 添加函数支持
找到torch安装位置,打开torch/onnx/symbolic_opset18.py,添加下列代码:
@_onnx_symbolic("aten::bitwise_right_shift")
@_beartype.beartype
def bitwise_right_shift(g, input: torch.Value, other: torch.Value):
return g.op("BitShift", input, other, direction_s='RIGHT')
@_onnx_symbolic("aten::bitwise_and")
@_beartype.beartype
def bitwise_and(g, input: torch.Value, other: torch.Value):
return g.op("BitwiseAnd", input, other)
4. 运行
可能会遇到一些数据类型的报错,根据ONNX文档修改torch运行过程中涉及到tensor的dtype即可.