Python 多态加载 ONNX 模型类 | 简化模型管理

以下是一个简单的 Python 代码实现，使用多态来加载不同的 ONNX 模型：

import onnxruntime as ort

class ONNXModel:
    '定义加载 ONNX 模型的基类'
    def __init__(self, model_path):
        self.session = ort.InferenceSession(model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
    
    def predict(self, input_data):
        '使用加载的模型进行预测'
        input_data = input_data.astype('float32')
        output = self.session.run([self.output_name], {self.input_name: input_data})
        return output[0]
    
class ModelA(ONNXModel):
    '加载 ModelA 的 ONNX 模型'
    def __init__(self):
        super().__init__('model_a.onnx')

class ModelB(ONNXModel):
    '加载 ModelB 的 ONNX 模型'
    def __init__(self):
        super().__init__('model_b.onnx')

在这个示例中，我们首先定义了基类 ONNXModel，它包含了加载模型和预测的基本方法。然后创建了两个子类 ModelA 和 ModelB，分别继承 ONNXModel 并指定要加载的 ONNX 模型文件。

利用多态，我们可以根据需求选择不同的模型类：

model = ModelA()  # 加载 ModelA 的 ONNX 模型
output = model.predict(input_data)  # 使用 ModelA 进行预测

model = ModelB()  # 加载 ModelB 的 ONNX 模型
output = model.predict(input_data)  # 使用 ModelB 进行预测

注意:

• input_data 应为一个 NumPy 数组，其形状和数据类型需要与模型输入节点匹配。
• onnxruntime 支持 CPU 和 GPU 加速，可以使用 ort.InferenceSession 的 providers 参数选择计算设备，例如：

session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])

• 在使用 GPU 加速时，需要确保 CUDA 已安装并且支持当前的 GPU，否则会引发运行时错误。

通过这种方式，我们可以方便地管理和加载不同的 ONNX 模型，并使用多态特性简化模型选择和预测过程。