如何在ONNX中使用TopK算子？自定义Sklearn转换器实践

我正在编写转换器和计算器，将自定义Sklearn转换器转换为ONNX模型，需要计算数据点的中位数。发现ONNX没有直接的中位数计算函数，因此用TopK算子实现，但TopK会返回两个输出（Top K值及对应索引），使用时遇到了问题：

• 首次尝试仅指定一个输出名，报错：

ValueError: Unexpected index 1 in operator name 'TopK' with .output names ['variable']

• 修改为指定两个输出名后，又因图断开报错：

RuntimeError: After 2 iterations for 2 nodes, still unable to sort names {'variable'}...

我只需要Top K的values而非indices，且要匹配算子输出要求。试过用ReduceSum包裹TopK的第一个输出可以运行，但不符合需求，最终用Identity算子包裹TopK[0]实现了预期效果：

Y = OnnxIdentity(OnnxTopK(X, np.array([3]),op_version=opv)[0], op_version=opv, output_names=out[:1])

请问有没有更简洁直接的实现方式？

完整最小可运行示例（MWE） #

import numpy as np
import pandas as pd

from onnxruntime import InferenceSession

from sklearn.base import BaseEstimator, TransformerMixin
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType, DoubleTensorType, guess_numpy_type
from skl2onnx.algebra.onnx_ops import (
    OnnxReduceSum,
    OnnxTopK,
    OnnxIdentity
)
from skl2onnx import update_registered_converter

def mt_transformer_shape_calculator(operator):
    op = operator.raw_operator
    input_type = operator.inputs[0].type.__class__
    input_dim = operator.inputs[0].get_first_dimension()
    n = operator.inputs[0].get_second_dimension()
    
    output_type = input_type([input_dim, 3])
    operator.outputs[0].type = output_type
    
def mt_transformer_converter(scope, operator, container):
    op = operator.raw_operator
    opv = container.target_opset
    out = operator.outputs

    X = operator.inputs[0]
    n = operator.inputs[0].get_second_dimension()
    dtype = guess_numpy_type(X.type)
    Y = OnnxIdentity(OnnxTopK(X, np.array([3]),op_version=opv)[0], op_version=opv, output_names=out[:1])
    Y.add_to(scope, container)
    

class MedianTransformer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        return self

    def transform(self, X):
        pass



data = pd.DataFrame(
    [[1,2,3,4],[4,5,6,5]]
)
    
update_registered_converter(
    MedianTransformer, "MTTransformer",
    mt_transformer_shape_calculator,
    mt_transformer_converter)

mt = MedianTransformer()
onx = convert_sklearn(mt, name='test', initial_types=[("X", FloatTensorType([None,4]))], 
                      final_types=[("Y", DoubleTensorType([None,3]))])


sess = InferenceSession(onx.SerializeToString())

sess.run(None, {'X': data.values.astype(np.float32)})[0]

输出 #

array([[4., 3., 2.],
       [6., 5., 5.]])

有两种更简洁的实现方式，不需要借助OnnxIdentity：

1. 直接指定TopK的输出名列表，忽略第二个输出
   在创建OnnxTopK时，直接传入包含两个输出名的列表，其中第二个输出名可以用一个临时名称（后续不会用到），这样既满足算子的输出要求，又能只保留需要的values输出。修改转换器代码如下：

def mt_transformer_converter(scope, operator, container):
    op = operator.raw_operator
    opv = container.target_opset
    out = operator.outputs

    X = operator.inputs[0]
    # 生成一个临时名称用于索引输出
    temp_idx = scope.get_unique_variable_name("topk_idx")
    # 直接指定两个输出，第一个是目标输出，第二个是临时索引
    OnnxTopK(X, np.array([3]), op_version=opv, output_names=[out[0].full_name, temp_idx]).add_to(scope, container)

2. 直接绑定TopK的第一个输出到目标节点
   利用OnnxTopK的输出索引选择，直接将第一个输出绑定到目标输出，skl2onnx会自动处理临时变量的管理，不需要额外算子：

def mt_transformer_converter(scope, operator, container):
    op = operator.raw_operator
    opv = container.target_opset
    out = operator.outputs

    X = operator.inputs[0]
    topk_node = OnnxTopK(X, np.array([3]), op_version=opv)
    # 直接将topk的第一个输出绑定到目标输出
    topk_node[0].add_to(scope, container, output_names=out[:1])

这两种方式都能达到和OnnxIdentity相同的效果，且更简洁，不需要额外的算子节点。

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