面向大规模部署的特性¶
本说明讨论了在大型系统中运行 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;
});
面向开发者的说明：可以使用 C10_LOG_API_USAGE_ONCE("my_api") 在 C++ 代码中添加新的 API 触发点，或使用 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库直接集成。可以通过扩展
Dataset或IterableDataset来利用torch.utils.data的现有功能。
