大规模部署的功能¶
本笔记讨论了几个扩展点和技巧,这些可能在将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;
});
开发者注意事项:新的API触发点可以在代码中添加,使用C++中的
C10_LOG_API_USAGE_ONCE("my_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库来实现。可以扩展
torch.utils.data的现有功能,或者使用Dataset或IterableDataset。