【Triton inference server】torch模型配置

Triton inference server中的torch模型配置

通常来说，使用Triton的时候更推荐使用Onnx模型进行推理，容错性更高。若想使用torch模型则需要将模型转换为torchscript格式。

• 模型配置文件示例

  name: "model_torchscript"
  platform: "pytorch_libtorch"
  max_batch_size : 0
  input [
    {
      name: "INPUT__0"
      data_type: TYPE_FP32
      dims: [1,1,3,8,224,224]
      
    }
  ]
  output [
    {
      name: "OUTPUT__0"
      data_type: TYPE_FP32
      dims: [ 1, 400 ]
    }
  ]
  # instance_group [
  #   {
  #     kind: KIND_GPU
  #     gpus: [ 0 ]
  #   }
  # ]

在配置文件中，默认的输入值是INPUT__0，对应torch代码中torch.nn.Module模块下的def forward(self, img)函数中的img参数。同理OUTPUT__0对应的是模型的第一个输出值。
若模型有多个输出则对应的参考配置文件如下:

• 多输出配置文件示例

name: "model_pt"
platform: "pytorch_libtorch"
max_batch_size : 0
input [
  {
    name: "INPUT__0"
    data_type: TYPE_FP32
    dims: [1, 3, 550, 550]
    
  }
]
output [
  {
    name: "OUTPUT__0"
    data_type: TYPE_FP32
    dims: [100, 4]
  },
  {
    name: "OUTPUT__1"
    data_type: TYPE_FP32
    dims: [100, 32]
  },
  {
    name: "OUTPUT__2"
    data_type: TYPE_INT64
    dims: [100]
  },
  {
    name: "OUTPUT__3"
    data_type: TYPE_FP32
    dims: [100]
  },
  {
    name: "OUTPUT__4"
    data_type: TYPE_FP32
    dims: [138, 138, 32]
  }
]

需要注意的是，配置文件中的data_type需要根据输出的数据类型进行设置。