• Docs >
  • Features for large-scale deployments
Shortcuts

大规模部署的功能

本笔记讨论了在大型系统中运行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 操作符运行时调用。通常,在初始化期间调用它们是一个好主意。

这是一个例子:

// 在程序开始时调用
void init() {
    // 随机采样百分之一的操作符运行
    addGlobalCallback(
      RecordFunctionCallback(
        &onFunctionEnter,
        &onFunctionExit)
      .needsInputs(true)
      .samplingProb(0.01)
    );
    // 注意,要在模型调用线程中启用观察器,
    // 在运行模型之前在线程中调用enableRecordFunction()
}

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被使用: " << event_name << std::endl;
});

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

将元数据附加到保存的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的现有功能,通过扩展DatasetIterableDataset来实现。