大规模部署功能¶
本说明介绍了在大型系统中运行 PyTorch 或在大型组织中使用 PyTorch 操作多个系统时可能有用的一些扩展点和技巧。
它不涵盖将模型部署到生产环境的主题。请查看 torch.jit
或相应的教程之一。
本说明假设您要么在组织中从源代码构建 PyTorch,要么能够静态链接在使用 PyTorch 时加载的其他代码。因此,许多钩子都以 C++ API 的形式公开,这些 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;
});
开发人员注意事项:可以使用 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 通常是松散耦合的,并且很容易用专门的版本替换组件。常见的扩展点包括
用 C++ 实现的自定义算子 - 请参阅 教程了解更多详情。
自定义数据读取通常可以通过调用相应的 Python 库直接集成。可以通过扩展
torch.utils.data
的现有功能来利用Dataset
或IterableDataset
。