推理模式¶
c10::InferenceMode
是一个新的 RAII 保护机制,类似于 NoGradMode
,当您确定您的操作不会与 autograd(例如模型训练)交互时使用。与 NoGradMode
相比,在此模式下运行的代码通过禁用与 autograd 相关的工作(如视图跟踪和版本计数器递增)获得更好的性能。但是,在 c10::InferenceMode
内部创建的张量在与 autograd 系统交互时也具有更多限制。
InferenceMode
可以为给定的代码块启用。在 InferenceMode
内部,所有新分配的(非视图)张量都标记为推理张量。推理张量
没有版本计数器,因此如果您尝试读取其版本(例如,因为您保存了此张量用于反向传播),则会引发错误。
在
InferenceMode
外部是不可变的。因此,如果您尝试执行以下操作,则会引发错误: - 在 InferenceMode 外部更改其数据。 - 在 InferenceMode 外部将其更改为requires_grad=True
。为了解决这个问题,您可以在InferenceMode
外部克隆以获得正常张量,然后再进行更改。
当且仅当非视图张量在 InferenceMode
内部分配时,它才是推理张量。当且仅当视图张量所视图的张量是推理张量时,视图张量才是推理张量。
在 InferenceMode
代码块内部,我们做出以下性能保证
像
NoGradMode
一样,即使操作的输入具有requires_grad=True
,所有操作也不会记录grad_fn
。这适用于推理张量和普通张量。对推理张量的视图操作不进行视图跟踪。视图和非视图推理张量是无法区分的。
保证对推理张量的原地操作不会执行版本递增。
有关 InferenceMode
的更多实现细节,请参阅 RFC-0011-InferenceMode。
从 AutoNonVariableTypeMode
迁移指南¶
在 PyTorch 的生产环境中使用推理工作负载时,我们已经看到 C++ 保护机制 AutoNonVariableTypeMode
(现在是 AutoDispatchBelowADInplaceOrView
)的广泛使用,它禁用了 autograd、视图跟踪和版本计数器递增。不幸的是,当前对推理工作负载使用此保护机制的习惯用法是不安全的:可以使用 AutoNonVariableTypeMode
绕过 PyTorch 的安全检查,并导致静默的错误结果,例如,当为反向传播保存的张量随后被更改时,PyTorch 会抛出错误,但更改发生在 AutoNonVariableTypeMode
内部会静默绕过检查,并向用户返回错误的梯度。
当 AutoNonVariableTypeMode
的当前用户考虑迁移时,以下步骤可能有助于您决定最佳替代方案
尝试在仅推理模式下运行工作负载的用户(例如加载预训练的 JIT 模型并在 C++ 运行时中运行推理)应添加
c10::InferenceMode guard
以保护张量上的所有操作(包括模型加载)。请参阅下面的推理工作负载示例
c10::InferenceMode guard;
model.load_jit(saved_model);
auto inputs = preprocess_tensors(data);
auto out = model.forward(inputs);
auto outputs = postprocess_tensors(out);
请注意,c10::InferenceMode
为 AutoNonVariableTypeMode
提供了直接替换方案,该方案保留了 AutoNonVariableTypeMode
的性能特征。但是它们也有一些用户应额外注意的差异
两种保护机制都会影响张量执行过程,以跳过与推理无关的工作,但
InferenceMode
也会影响张量创建,而AutoNonVariableTypeMode
则不会。换句话说,在InferenceMode
内部创建的张量被标记为推理张量,以便在退出InferenceMode
后可以应用某些限制。启用/禁用
InferenceMode
状态可以嵌套,而AutoNonVariableTypeMode
仅允许启用状态。
{
InferenceMode guard(true);
// InferenceMode is on
{
InferenceMode guard(false);
// InferenceMode is off
}
// InferenceMode is on
}
// InferenceMode is off
尝试实现自定义内核且想要在
Autograd
调度键下重新调度的用户应使用AutoDispatchBelowADInplaceOrView
代替。请注意,AutoDispatchBelowADInplaceOrView
只是AutoNonVariableTypeMode
的新名称,因为它更好地解释了保护机制的功能。我们正在弃用AutoNonVariableTypeMode
,它将在 1.10 版本中删除。有关示例,请参阅pytorch/vision
中的自定义内核ROIAlignFunction
class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
public:
static torch::autograd::variable_list forward(
torch::autograd::AutogradContext* ctx,
const torch::autograd::Variable& input,
const torch::autograd::Variable& rois,
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width,
int64_t sampling_ratio,
bool aligned) {
ctx->saved_data["spatial_scale"] = spatial_scale;
ctx->saved_data["pooled_height"] = pooled_height;
ctx->saved_data["pooled_width"] = pooled_width;
ctx->saved_data["sampling_ratio"] = sampling_ratio;
ctx->saved_data["aligned"] = aligned;
ctx->saved_data["input_shape"] = input.sizes();
ctx->save_for_backward({rois});
// Used to be at::AutoNonVariableTypeMode g;
at::AutoDispatchBelowADInplaceOrView guard;
auto result = roi_align(
input, rois, spatial_scale, pooled_height,
pooled_width, sampling_ratio, aligned);
return {result};
}
自定义原地和视图内核除了上述保护机制外,还需要一些特殊处理,有关更多详细信息,请参阅自定义内核教程。