推断模式#
c10::InferenceMode 是一种新的 RAII 守护(guard),类似于 NoGradMode,适用于您确定操作不会与 autograd 交互的情况(例如模型训练)。与 NoGradMode 相比,在此模式下运行的代码通过禁用 autograd 相关的操作(如视图跟踪和版本计数器递增)来获得更好的性能。但是,在 c10::InferenceMode 中创建的张量在与 autograd 系统交互时也有更多限制。
InferenceMode 可以为给定的代码块启用。在 InferenceMode 内部,所有新分配的(非视图)张量都会被标记为推断张量。推断张量
没有版本计数器,因此如果您尝试读取它们的版本(例如,因为您保存了此张量以用于反向传播),将会引发错误。
在
InferenceMode外部是不可变的。因此,如果您尝试:- 在 InferenceMode 外部修改其数据。- 在 InferenceMode 外部将其修改为requires_grad=True。将会引发错误。为了解决这个问题,您可以在 InferenceMode 外部创建一个克隆,以便在修改之前获得一个常规张量。
当且仅当非视图张量在 InferenceMode 内部分配时,它才是推断张量。当且仅当视图张量是推断张量的视图时,它才是推断张量。
在 InferenceMode 块内,我们提供了以下性能保证:
与
NoGradMode类似,所有操作都不会记录grad_fn,即使它们的输入具有requires_grad=True。这适用于推断张量和常规张量。对推断张量进行视图操作不会进行视图跟踪。视图和非视图推断张量是无法区分的。
对推断张量进行原地(inplace)操作保证不会进行版本递增。
有关 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
尝试实现自定义内核并希望在
Autograddispatch keys 下重新分派的用户应改用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};
}
自定义原地(inplace)和视图(view)内核除了上述守护之外,还需要一些特殊处理,有关更多详细信息,请参阅 自定义内核教程。