model.py 831 Bytes
Newer Older
root's avatar
root committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import triton_python_backend_utils as pb_utils

class TritonPythonModel:
    def initialize(self, args):
        """模型初始化时调用(可选)"""
        print("initialized")


    def execute(self, requests):
        responses = []
        for request in requests:
            #获取输入参数
            in_0 = pb_utils.get_input_tensor_by_name(request, "INPUT0").as_numpy()
            in_1 = pb_utils.get_input_tensor_by_name(request, "INPUT1").as_numpy()
            
            # 逻辑计算
            out_tensor = pb_utils.Tensor("OUTPUT0", in_0 + in_1)
            
            # 封装并返回
            responses.append(pb_utils.InferenceResponse([out_tensor]))
        return responses

    def finalize(self):
            """模型卸载时调用(可选)"""
            print('Cleaning up...')