准备

创建模型仓库

mkdir torch_test
cd torch_test/
mkdir -p ./model_repository/fc_model_onnx/1
mkdir -p ./model_repository/fc_model_pt/1

为模型库生成2个测试模型 （torch和onnx）

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        # 嵌入层
        self.embedding = nn.Embedding(100, 8)
        # 全连接层
        self.fc = nn.Linear(8, 4)
        # 包含了4个全连接层的序列模块（nn.Sequential）
        self.fc_list = nn.Sequential(*[nn.Linear(8, 8) for _ in range(4)])

    def forward(self, input_ids):
        # 输入 input_ids 经过嵌入层后得到词向量 word_emb
        word_emb = self.embedding(input_ids)
        output1 = self.fc(word_emb)
        output2 = self.fc_list(word_emb)
        return output1, output2

if __name__ == "__main__":
    pt_path = "./model_repository/fc_model_pt/1/model.pt"
    onnx_path = "./model_repository/fc_model_onnx/1/model.onnx"
    model = SimpleModel()
    # 创建示例输入,该输入应与模型的实际输入具有相同的形状和数据类型
    ipt = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
    # script_model 将成为一个被优化和编译的模型表示，并且可以被保存和重新加载以供后续使用。
    script_model = torch.jit.trace(model, ipt, strict=True)
    # script_model 保存为PyTorch模型文件，文件路径为 pt_path
    torch.jit.save(script_model, pt_path)
    # torch.onnx.export() 函数将模型 model 导出为ONNX格式的模型文件，文件路径为 onnx_path。同时指定了输入名字为 'input'，输出名字分别为 'output1' 和 'output2'。
    torch.onnx.export(model, ipt, onnx_path,
                      input_names=['input'],
                      output_names=['output1', 'output2']
                      )

jit.trace介绍

torch.jit.trace 是 PyTorch 中的一个函数，用于将给定的模型或模块进行即时编译（Just-in-Time Compilation，JIT）并追踪其执行过程。通过追踪模型的执行，可以将其转换为一个高性能的图形计算表示，以便在后续的推理过程中提供更快的执行速度。

生成模型

pip3 install onnx -i ****s://pypi.tuna.tsinghua.edu.cn/simple
python3 generate_testM.py

在"./model_repository/fc_model_pt" 目录下新建 config.pbtxt 文件

name: "fc_model_pt" # 模型名，也是目录名
platform: "pytorch_libtorch" # torch模型对应的平台
max_batch_size : 64 # 一次送入模型的最大batch_size
input [
  {
    name: "input_0" 
    data_type: TYPE_INT64 
    dims: [ -1 ]  # -1 代表是可变维度,输入是二维(默认第一个参数是batch_size)
  }
]
output [
  {
    name: "output_0" 
    data_type: TYPE_FP32
    dims: [ -1, -1, 4 ]
  },
  {
    name: "output_1"
    data_type: TYPE_FP32
    dims: [ -1, -1, 8 ]
  }
]

启动Triton服务

docker run --gpus=1 --rm --net=host  -v ./model_repository:/models tritonserver:23.05-py3 tritonserver --model-repository=/models --strict-model-config=false
# 上面命令中 strict-model-config=false，这个表示使用Triton自动生成模型的配置
# 如果熟悉模型配置的规范，可以先自己配置好config文件。再启动Tritont服务，对应的strict-model-config传True
#  '--strict-model-config' 的参数，用于启用严格的模型配置检查。然而，该参数已经被认为过时，不再建议使用。
# --disable-auto-complete-config
# 作为替代，新的参数 '--disable-auto-complete-config' 被引入，用于禁用 Triton Inference Server 的自动模型配置补全功能。该功能在旧版本中默认启用，它可以自动补全缺失的模型配置，但有时也会导致意外行为。
# 不加参数默认自动生成config文件

使用命令查看模型生成的配置文件

# curl localhost:8000/v2/models/<your model name>/config 
curl localhost:8000/v2/models/fc_model_pt/config

得到输出配置后，可以自己新建onnx模型的config.pbtxt文件

name: "fc_model_onnx" 
platform: "onnxruntime_onnx" 
max_batch_size : 0 
input [
  {
    name: "input" 
    data_type: TYPE_INT64 
    dims: [2, 3]  
  }
]
output [
  {
    name: "output1" 
    data_type: TYPE_FP32
    dims: [2, 3, 4]
  },
  {
    name: "output2"
    data_type: TYPE_FP32
    dims: [2, 3, 8]
  }
]

有了自定义配置文件后，可以重启triton

docker run --gpus=1 --rm --net=host  -v ./model_repository:/models tritonserver:23.05-py3 tritonserver --model-repository=/models --strict-model-config=True

****测试

import numpy as np
import tritonclient.**** as ****client

triton_client = ****client.InferenceServerClient(url="localhost:8000", verbose=False)

model_name = "fc_model_pt"

inputs = [
    ****client.InferInput('input_0', [2, 3], "INT64")
]
outputs = [
    ****client.InferRequestedOutput('output_0'),
    ****client.InferRequestedOutput('output_1')
]

inputs[0].set_data_from_numpy(np.random.randint(0, high=5, size=(2, 3)))

results = triton_client.infer(model_name=model_name, inputs=inputs, outputs=outputs)

print(results.as_numpy("output_0"))
print(results.as_numpy("output_1"))

GRPC测试

import numpy as np
import tritonclient.grpc as grpcclient

triton_client = grpcclient.InferenceServerClient(url="localhost:8001", verbose=False)

model_name = "fc_model_pt"

inputs = [
    grpcclient.InferInput('input_0', [2, 3], "INT64")
]
outputs = [
    grpcclient.InferRequestedOutput('output_0'),
    grpcclient.InferRequestedOutput('output_1')
]

inputs[0].set_data_from_numpy(np.random.randint(0, high=5, size=(2, 3)))

results = triton_client.infer(model_name=model_name, inputs=inputs, outputs=outputs)

print(results.as_numpy("output_0"))
print(results.as_numpy("output_1"))