Pytorch中的INFERENCE MODE

imoldpan · 2023 年7 月 5 日 07:37

由GPT4生产

这个链接是 PyTorch 的官方文档，主要介绍了一个新的 RAII 守卫 c10::InferenceMode，用于在你确定你的操作不会与自动梯度（autograd）有任何交互的情况下使用，例如模型推理。以下是一些重点：

InferenceMode 的使用：c10::InferenceMode 可以在给定的代码块中启用。在 InferenceMode 内部，所有新分配的（非视图）张量都被标记为推理张量。
推理张量的特性：推理张量没有版本计数器，所以如果你试图读取它们的版本（例如，因为你保存了这个张量用于反向传播），将会引发错误。推理张量在 InferenceMode 外部是不可变的，所以如果你试图在 InferenceMode 外部改变它们的数据或将它们变为 requires_grad=True，将会引发错误。为了解决这个问题，你可以在 InferenceMode 外部克隆一个普通的张量，然后进行修改。
InferenceMode 的性能保证：在 InferenceMode 块内，所有操作即使它们的输入有 requires_grad=True，也不会记录 grad_fn。这适用于推理张量和普通张量。对推理张量的视图操作不进行视图跟踪。推理张量的就地操作保证不会进行版本 bump。
从 AutoNonVariableTypeMode 迁移：对于在 PyTorch 中用于推理工作负载的生产使用，我们已经看到了 C++ 守卫 AutoNonVariableTypeMode（现在是 AutoDispatchBelowADInplaceOrView）的使用，它禁用了自动梯度、视图跟踪和版本计数器 bump。但是，当前这个守卫对于推理工作负载的使用是不安全的：使用 AutoNonVariableTypeMode 可能会绕过 PyTorch 的安全检查，导致错误的结果。当 AutoNonVariableTypeMode 的当前用户考虑迁移时，可以考虑使用 c10::InferenceMode 守卫来保护所有的张量操作（包括模型加载）。
InferenceMode 和 AutoNonVariableTypeMode 的区别：两者都影响张量执行过程，跳过与推理无关的工作，但是 InferenceMode 也影响张量创建，而 AutoNonVariableTypeMode 不影响。换句话说，在 InferenceMode 内部创建的张量被标记为推理张量，以便在退出 InferenceMode 后可以应用某些限制。InferenceMode 的启用/禁用状态可以嵌套，而 AutoNonVariableTypeMode 只允许启用状态。
自定义内核的实现：试图实现自定义内核并希望在 Autograd 分发键下重新分发的用户应该使用 AutoDispatchBelowADInplaceOrView。注意 AutoDispatchBelowADInplaceOrView 只是 AutoNonVariableTypeMode 的新名称，因为它更好地解释了守卫的功能。我们正在弃用 AutoNonVariableTypeMode，并将在 1.10 版本中移除它。

c10::InferenceMode 是一个用于在你确定你的操作不会与自动梯度（autograd）有任何交互的情况下使用的 RAII 守卫。你可以在给定的代码块中启用它。在 InferenceMode 内部，所有新分配的（非视图）张量都被标记为推理张量。

以下是一个使用 c10::InferenceMode 的代码示例：

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 守卫。然后，我们加载了一个 JIT 模型，对数据进行预处理，执行模型的前向传播，最后对输出进行后处理。所有这些操作都在 c10::InferenceMode 守卫的作用范围内进行，这意味着它们都被视为推理操作，不会与 autograd 有任何交互。

需要注意的是，c10::InferenceMode 的启用/禁用状态可以嵌套。例如：

{
  c10::InferenceMode guard(true);
  // InferenceMode is on
  {
    c10::InferenceMode guard(false);
    // InferenceMode is off
  }
  // InferenceMode is on
}
// InferenceMode is off

在这个示例中，我们首先启用了 InferenceMode，然后在内部的一个代码块中禁用了它，然后在外部的代码块中又重新启用了它。这种嵌套的启用/禁用状态可以让你更灵活地控制哪些代码应该在 InferenceMode 下运行，哪些代码不应该。

Pytorch中的INFERENCE MODE

参考