給本頁評分

推理模式#

c10::InferenceMode 是一種新的 RAII guard,類似於 NoGradMode,當您確定您的操作不會與 autograd 互動時(例如模型訓練),可以使用它。與 NoGradMode 相比,在此模式下執行的程式碼透過停用與 autograd 相關的工作(如檢視跟蹤和版本計數器遞增)來獲得更好的效能。然而,在 c10::InferenceMode 內建立的 tensors 在與 autograd 系統互動時也有更多限制。

可以為給定的程式碼塊啟用 InferenceMode。在 InferenceMode 內部,所有新分配的(非檢視)tensors 都被標記為推理 tensors。推理 tensors

  • 沒有版本計數器,因此如果您嘗試讀取它們的版本(例如,因為您儲存了此 tensor 用於反向傳播),將會引發錯誤。

  • InferenceMode 之外是不可變的。因此,如果您嘗試: - 在 InferenceMode 之外修改它們的資料。 - 在 InferenceMode 之外將它們修改為 requires_grad=True,將會引發錯誤。要解決此問題,您可以在 InferenceMode 之外進行克隆以獲得普通 tensor,然後再進行修改。

一個非檢視 tensor 是推理 tensor 當且僅當它是在 InferenceMode 內分配的。一個檢視 tensor 是推理 tensor 當且僅當它是推理 tensor 的檢視。

InferenceMode 塊內部,我們提供以下效能保證:

  • NoGradMode 一樣,即使輸入具有 requires_grad=True,所有操作也不會記錄 grad_fn。這適用於推理 tensors 和普通 tensors。

  • 對推理 tensors 的檢視操作不會進行檢視跟蹤。檢視推理 tensors 和非檢視推理 tensors 是無法區分的。

  • 對推理 tensors 的原地(inplace)操作保證不會進行版本遞增。

有關 InferenceMode 的更多實現細節,請參閱 RFC-0011-InferenceMode

AutoNonVariableTypeMode 的遷移指南#

在 PyTorch 用於推理工作負載的生產實踐中,我們看到 C++ guard AutoNonVariableTypeMode(現為 AutoDispatchBelowADInplaceOrView)被廣泛使用,它停用了 autograd、檢視跟蹤和版本計數器遞增。不幸的是,目前這種 guard 在推理工作負載中的普遍用法是不安全的:使用 AutoNonVariableTypeMode 可能會繞過 PyTorch 的安全檢查並導致靜默的錯誤結果,例如,當用於反向傳播的 tensors 隨後被修改時,PyTorch 會丟擲錯誤,但發生在 AutoNonVariableTypeMode 內的修改會靜默地繞過檢查並向用戶返回錯誤的梯度。

當前 AutoNonVariableTypeMode 使用者在考慮遷移時,以下步驟可能有助於您確定最佳替代方案:

  1. 嘗試僅在推理模式下執行工作負載(例如載入預訓練的 JIT 模型並在 C++ 執行時中執行推理)的使用者應新增 c10::InferenceMode guard 來保護所有 tensor 操作(包括模型載入)。請參閱下面的推理工作負載示例

c10::InferenceMode guard;
model.load_jit(saved_model);
auto inputs = preprocess_tensors(data);
auto out = model.forward(inputs);
auto outputs = postprocess_tensors(out);

注意,c10::InferenceModeAutoNonVariableTypeMode 提供了一個即時替換方案,保留了 AutoNonVariableTypeMode 的效能特性。但它們也有一些使用者需要額外注意的區別:

  • 這兩種 guard 都會影響 tensor 執行過程,跳過與推理無關的工作,但 InferenceMode 也會影響 tensor 的建立,而 AutoNonVariableTypeMode 不會。換句話說,在 InferenceMode 內部建立的 tensors 被標記為推理 tensors,以便在退出 InferenceMode 後應用某些限制。

  • 啟用/停用 InferenceMode 狀態可以巢狀,而 AutoNonVariableTypeMode 只允許啟用狀態。

{
  InferenceMode guard(true);
  // InferenceMode is on
  {
    InferenceMode guard(false);
    // InferenceMode is off
  }
  // InferenceMode is on
}
// InferenceMode is off
  1. 嘗試實現自定義 kernel,希望在 Autograd dispatch keys 下重新分派的使用者應改用 AutoDispatchBelowADInplaceOrView。注意,AutoDispatchBelowADInplaceOrView 只是 AutoNonVariableTypeMode 的新名稱,因為它更好地解釋了該 guard 的功能。我們正在棄用 AutoNonVariableTypeMode,它將在 1.10 版本中移除。請參閱 pytorch/vision 中自定義 kernel 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) 和檢視 kernel 除了上述 guard 外還需要一些特殊處理,詳情請參閱自定義 kernel 教程