大规模部署功能 ¶

本笔记讨论了在更大系统内运行 PyTorch 或在一个更大组织内使用 PyTorch 操作多个系统时可能有用的几个扩展点和技巧。

这部分内容没有涵盖将模型部署到生产环境的话题。请查看 torch.jit 或相应的教程之一。

该说明假设您要么在您的组织中从源代码构建 PyTorch，要么有在 PyTorch 使用时静态链接额外代码的能力。因此，许多钩子都作为 C++ API 暴露，可以在集中位置触发一次，例如在静态初始化代码中。

全集群操作员分析¶

PyTorch 自带 torch.autograd.profiler ，可以按需测量单个操作符所需的时间。可以使用相同的机制对任何运行 PyTorch 的进程进行“始终开启”的测量。这可能有助于收集有关给定进程中运行的 PyTorch 工作负载或整个机器集上的信息。

可以通过 torch::addGlobalCallback 添加任何操作符调用的新回调。钩子将使用 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 密钥”，用于标识使用点，例如 python.import 用于 PyTorch 扩展导入，或者 torch.script.compile 如果触发了 TorchScript 编译。

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

开发者笔记：新 API 触发点可以通过在 C++中使用 C10_LOG_API_USAGE_ONCE("my_api") 或在 Python 中使用 torch._C._log_api_usage_once("my.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 库进行集成。可以通过扩展 Dataset 或 IterableDataset 来利用 torch.utils.data 的现有功能。