大规模部署特性

本注意事项讨论了在较大系统内运行 PyTorch 或在大型组织中使用 PyTorch 操作多个系统时可能有用的一些扩展点和技巧。

本注意事项不涉及模型投入生产部署的主题。请查看 torch.jit 或相应的教程之一。

本注意事项假设您要么在组织内部从源代码构建 PyTorch，要么能够在 PyTorch 使用时静态链接要加载的附加代码。因此，许多钩子都作为 C++ API 公开，可以在集中位置触发一次，例如在静态初始化代码中。

全舰队（Fleet-wide）算子分析

PyTorch 提供了 torch.autograd.profiler，能够按需测量单个算子花费的时间。可以使用相同的机制对运行 PyTorch 的任何进程进行“始终开启”的测量。这对于收集有关给定进程或整个机器集群中运行的 PyTorch 工作负载的信息非常有用。

可以使用 torch::addGlobalCallback 为任何算子调用添加新的回调函数。钩子将通过描述调用上下文（例如 name）的 torch::RecordFunction 结构体进行调用。如果启用，RecordFunction::inputs() 包含表示为 torch::IValue 变体类型的函数参数。请注意，输入日志记录相对昂贵，因此必须显式启用。

算子回调函数还可以访问 c10::ThreadLocalDebugInfo::get() 接口，该接口返回指向保存调试信息的结构体的指针。此调试信息可以使用 at::DebugInfoGuard 对象提前设置。调试信息会通过前向（包括异步 fork 任务）和后向传播，对于将应用程序更高层的一些执行环境额外信息（例如模型 ID）传递给算子回调函数非常有用。

调用回调函数会增加一些开销，因此通常只随机抽样算子调用是有用的。这可以通过将可选的采样率传递给 torch::addGlobalCallback，针对每个回调函数启用。

请注意，addGlobalCallback 不是线程安全的，只能在没有 PyTorch 算子运行时调用。通常，最好在初始化期间调用它们一次。

这是一个例子

// Called somewhere in the program beginning
void init() {
    // Sample one in a hundred operator runs randomly
    addGlobalCallback(
      RecordFunctionCallback(
        &onFunctionEnter,
        &onFunctionExit)
      .needsInputs(true)
      .samplingProb(0.01)
    );
    // Note, to enable observers in the model calling thread,
    // call enableRecordFunction() in the thread before running a model
}

void onFunctionEnter(const RecordFunction& fn) {
    std::cerr << "Before function " << fn.name()
              << " with " << fn.inputs().size() << " inputs" << std::endl;
}

void onFunctionExit(const RecordFunction& fn) {
    std::cerr << "After function " << fn.name();
}

API 使用情况日志记录

在更广泛的生态系统（例如托管作业调度程序）中运行时，跟踪哪些二进制文件调用了特定的 PyTorch API 通常很有用。在几个重要的 API 点注入了简单的检测工具，可以触发给定的回调函数。由于 PyTorch 通常在一次性 Python 脚本中调用，因此对于每个 API，回调函数在给定进程中仅触发一次。

可以使用 c10::SetAPIUsageHandler 注册 API 使用情况检测处理程序。传递的参数将是标识使用点的“api key”，例如，对于 PyTorch 扩展导入是 python.import，如果触发了 TorchScript 编译，则是 torch.script.compile。

SetAPIUsageLogger([](const std::string& event_name) {
    std::cerr << "API was used: " << event_name << std::endl;
});

开发者注意事项：可以在代码中使用 C++ 中的 C10_LOG_API_USAGE_ONCE("my_api") 或 Python 中的 torch._C._log_api_usage_once("my.api") 添加新的 API 触发点。

为保存的 TorchScript 模型附加元数据

TorchScript 模块可以保存为存档文件，该文件将序列化的参数和模块代码捆绑为 TorchScript（参见 torch.jit.save()）。将附加信息（例如模型生产者描述或辅助工件）与模型捆绑在一起通常很方便。

这可以通过将 _extra_files 参数传递给 torch.jit.save() 和 torch::jit::load 来实现，以便在保存过程中存储和检索任意二进制大对象。由于 TorchScript 文件是常规的 ZIP 存档，额外信息将作为常规文件存储在存档的 extra/ 目录中。

还有一个全局钩子，允许将额外文件附加到当前进程中生成的任何 TorchScript 存档。这对于使用生产者元数据标记模型非常有用，类似于数码相机生成的 JPEG 元数据。示例用法可能如下所示

SetExportModuleExtraFilesHook([](const Module&) {
    ExtraFilesMap files;
    files["producer_info.json"] = "{\"user\": \"" + getenv("USER") + "\"}";
    return files;
});

构建环境注意事项

TorchScript 的编译需要访问原始的 Python 文件，因为它使用了 Python 的 inspect.getsource 调用。在某些生产环境中，可能需要将 .py 文件与预编译的 .pyc 文件一起显式部署。

常见扩展点

PyTorch API 通常是松耦合的，很容易用专用版本替换组件。常见的扩展点包括

• 用 C++ 实现自定义算子 - 详情请参阅教程。

• 自定义数据读取通常可以通过调用相应的 Python 库直接集成。可以通过扩展 torch.utils.data 的现有功能来利用 Dataset 或 IterableDataset。