参数

• model (torch.nn.Module) — 我们想要为流水线并行推理拆分的模型
• split_points (str 或 List[str]，默认为 ‘auto’) — 如何生成分割点并在每个 GPU 上对模型进行分块。‘auto’ 将根据任何模型找到最佳的平衡分割。否则，应为模型中要按其分割的层名称列表。
• no_split_module_classes (List[str]) — 我们不希望被分割的层的类名列表。
• example_args (模型输入的元组) — 使用基于顺序的输入的模型的预期输入，用于单个进程。如果可能，建议使用此方法。
• example_kwargs (模型输入的字典) — 使用基于字典的输入的模型的预期输入，用于单个进程。这是一种非常有限的结构，要求在所有推理调用中都存在相同的键。除非所有情况下的先前条件都为真，否则不建议使用。
• num_chunks (int，默认为可用 GPU 的数量) — 流水线将具有的不同阶段的数量。默认情况下，它会为每个 GPU 分配一个块，但这可以进行调整和修改。通常，应该有 num_chunks >= num_gpus。
• gather_output (bool，默认为 False) — 如果为 True，则来自最后一个 GPU（包含真实输出）的输出将发送到所有 GPU。