• 文档 >
  • 大规模部署的特性
快捷方式

大规模部署的特性

本说明讨论了在更大的系统中使用 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 密钥”,例如 python.import 用于 PyTorch 扩展导入,或者 torch.script.compile 如果触发了 TorchScript 编译。

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,以在保存过程中存储和检索任意二进制 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 通常是松耦合的,并且很容易用专门的版本替换组件。常用扩展点包括

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源