Triton教程 — 动态批处理

Triton系列教程:

1. 快速开始
2. 利用Triton部署你自己的模型
3. Triton架构
4. 模型仓库
5. 存储代理
6. 模型设置
7. 优化
8. 动态批处理

Triton 提供了动态批处理功能，将多个请求组合在一起执行同一模型以提供更大的吞吐量。 默认情况下，只有当每个输入在请求中具有相同的形状时，请求才能被动态批处理。 为了在输入形状经常变化的情况下利用动态批处理，客户端需要将请求中的输入张量填充为相同的形状。

参差不齐的批处理是一种通过允许用户指定哪些输入不需要形状检查来避免显式填充的功能。 用户可以通过在模型配置中设置 allow_ragged_batch 字段来指定此类输入（参差不齐的输入）：

...
input [
  {
    name: "input0"
    data_type: TYPE_FP32
    dims: [ 16 ]
    allow_ragged_batch: true
  }
]
...

如何在一批请求中处理参差不齐的输入取决于后端实现。 ONNX Runtime 后端、TensorFlow 后端、PyTorch 后端和 TensorRT 后端等后端要求模型将参差不齐的输入作为一维张量接受。 这些后端将请求输入连接到一维张量中。

由于连接的输入不跟踪每个请求的开始和结束索引，因此后端通常要求模型具有额外的输入，批输入，描述有关形成的批的各种信息。

批量输入

批输入通常与参差不齐的输入结合使用，以提供有关每个批元素的信息，例如批中每个请求的输入元素计数。 批次输入由 Triton 生成，而不是在请求中提供，因为信息只有在动态批次形成后才能最终确定。

除了元素计数，用户还可以指定其他批输入类型，详情请参阅 protobuf 文档。

参差不齐的输入和批量输入的例子

如果您的模型接受 1 个可变长度输入张量 INPUT，形状为 [ -1, -1 ]。 第一个维度是批量维度，第二个维度是变长内容。 当客户端发送 3 个形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ] 的请求时。 为了利用动态批处理，实现此模型的直接方法是期望输入形状 [-1, -1] 并假设所有输入都被填充到相同的长度，以便所有请求都变成形状 [1, 5]，因此 Triton 可以 批处理并将它们作为单个 [ 3, 5 ] 张量发送到模型。 在这种情况下，填充张量和填充内容的额外模型计算将产生开销。 下面是输入配置：

max_batch_size: 16
input [
  {
    name: "INPUT"
    data_type: TYPE_FP32
    dims: [ -1 ]
  }
]

使用 triton 参差不齐的批处理，模型将实现为期望输入形状 [-1] 和一个额外的批输入、索引、形状 [-1]，模型应该使用它来解释输入中的批元素。 对于这样的模型，客户端请求不需要填充，可以按原样发送（形状为[1, 3]，[1, 4]，[1, 5]）。 上面讨论的后端会将输入批处理为形状为 [12] 的张量，其中包含请求的 3 + 4 + 5 级联。 Triton 还创建了形状为 [ 3 ] 且值为 [ 3, 7, 12 ] 的批量输入张量，它给出了每个批量元素结束的输入张量的偏移量。 下面是输入配置：

max_batch_size: 16
input [
  {
    name: "INPUT"
    data_type: TYPE_FP32
    dims: [ -1 ]
    allow_ragged_batch: true
  }
]
batch_input [
  {
    kind: BATCH_ACCUMULATED_ELEMENT_COUNT
    target_name: "INDEX"
    data_type: TYPE_FP32
    source_input: "INPUT"
  }
]

上面的示例使用 BATCH_ACCUMULATED_ELEMENT_COUNT 类型的参差不齐的批处理。 protobuf 文档中描述的其他类型的操作类似。