大规模部署功能¶
此说明介绍了几个扩展点和技巧,这些技巧在大型系统中运行 PyTorch 或在大型组织中使用 PyTorch 运行多个系统时可能有用。
它不涉及将模型部署到生产环境的主题。请查看 torch.jit 或相应的教程之一。
此说明假设您在组织中从源代码构建 PyTorch,或者能够静态链接其他代码,以便在使用 PyTorch 时加载这些代码。因此,许多挂钩都公开为 C++ API,可以在集中位置(例如静态初始化代码)中触发一次。
全集群运算符分析¶
PyTorch 附带 torch.autograd.profiler,能够按需测量各个运算符所花费的时间。人们可以使用相同的机制对运行 PyTorch 的任何进程执行“始终开启”测量。这可能有助于收集有关在给定进程或整个机器集中运行的 PyTorch 工作负载的信息。
可以使用 torch::addGlobalCallback 为任何运算符调用添加新的回调。将使用 torch::RecordFunction 结构调用挂钩,该结构描述调用上下文(例如 name)。如果启用,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 密钥”,例如 PyTorch 扩展导入的 python.import 或触发 TorchScript 编译的 torch.script.compile。
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 在保存过程中存储和检索任意二进制 blob 来实现。由于 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的现有功能。