InplaceFunction

class torch.autograd.function.InplaceFunction(inplace=False)

這個類別僅為向後相容性而存在。對於任何新的使用案例，請使用 Function 來取代它。

static backward(ctx, *grad_outputs)

定義一個使用反向模式自動微分來微分運算的公式。

這個函數需要被所有子類別覆寫。（定義這個函數等同於定義 vjp 函數。）

它必須接受一個 context ctx 作為第一個引數，接著是和 forward() 傳回的一樣多的輸出（對於 forward 函數的非張量輸出，將傳入 None），並且它應該傳回和 forward() 的輸入一樣多的張量。每個引數是相對於給定輸出的梯度，並且每個傳回的值應該是相對於相應輸入的梯度。如果輸入不是 Tensor，或者是不需要梯度的 Tensor，您可以直接將 None 作為該輸入的梯度傳入。

context 可用於檢索在 forward 傳遞期間儲存的張量。它還有一個屬性 ctx.needs_input_grad，它是一個布林值元組，表示每個輸入是否需要梯度。例如，如果 forward() 的第一個輸入需要計算相對於輸出的梯度，則 backward() 將具有 ctx.needs_input_grad[0] = True。

傳回型別

Any

static forward(*args, **kwargs)

定義自定義 autograd Function 的 forward 傳遞。

這個函數需要被所有子類別覆寫。有兩種方式來定義 forward 傳遞

用法 1（結合 forward 和 ctx）

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• 它必須接受一個 context ctx 作為第一個引數，接著是任意數量的引數（張量或其他型別）。

• 請參閱 結合或分離 forward() 和 setup_context() 以獲取更多詳細資訊

用法 2（分離 forward 和 ctx）

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• forward 不再接受 ctx 引數。

• 相反地，您還必須覆寫 torch.autograd.Function.setup_context() staticmethod 來處理設定 ctx 物件。output 是 forward 的輸出，inputs 是 forward 的輸入元組。

• 請參閱 擴充 torch.autograd 以獲取更多詳細資訊

context 可用於儲存任意資料，這些資料然後可以在 backward 傳遞期間檢索。張量不應直接儲存在 ctx 上（儘管目前為了向後相容性而沒有強制執行）。相反地，張量應該使用 ctx.save_for_backward() 儲存，如果它們打算在 backward（等效地，vjp）中使用，或者使用 ctx.save_for_forward() 儲存，如果它們打算在 jvp 中使用。

傳回型別

Any

static jvp(ctx, *grad_inputs)

定義一個使用前向模式自動微分來微分運算的公式。

這個函數需要被所有子類別覆寫。它必須接受一個 context ctx 作為第一個引數，接著是和 forward() 得到的輸入一樣多的輸入（對於 forward 函數的非張量輸入，將傳入 None），並且它應該傳回和 forward() 的輸出一樣多的張量。每個引數是相對於給定輸入的梯度，並且每個傳回的值應該是相對於相應輸出的梯度。如果輸出不是 Tensor，或者該函數對於該輸出不可微分，您可以直接將 None 作為該輸入的梯度傳入。

您可以使用 ctx 物件將任何值從 forward 傳遞到此函數。

傳回型別

Any

mark_dirty(*args)

將給定的張量標記為在原地運算中已修改。

這應該最多呼叫一次，無論是在 setup_context() 還是 forward() 方法中，並且所有引數都應該是輸入。

每次在呼叫 forward() 中原地修改的每個張量都應該提供給這個函數，以確保我們檢查的正確性。函數是在修改之前還是之後呼叫並不重要。

範例：

>>> class Inplace(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         x_npy = x.numpy() # x_npy shares storage with x
>>>         x_npy += 1
>>>         ctx.mark_dirty(x)
>>>         return x
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, grad_output):
>>>         return grad_output
>>>
>>> a = torch.tensor(1., requires_grad=True, dtype=torch.double).clone()
>>> b = a * a
>>> Inplace.apply(a)  # This would lead to wrong gradients!
>>>                   # but the engine would not know unless we mark_dirty
>>> b.backward() # RuntimeError: one of the variables needed for gradient
>>>              # computation has been modified by an inplace operation

mark_non_differentiable(*args)

將輸出標記為不可微分。

此方法最多應調用一次，無論是在 setup_context() 還是 forward() 方法中，並且所有參數都應該是 Tensor 輸出。

這會將輸出標記為不需要梯度，從而提高反向計算的效率。您仍然需要在 backward() 中為每個輸出接受一個梯度，但它始終會是一個零 Tensor，其形狀與相應輸出的形狀相同。

例如，這用於從排序返回的索引。請參閱範例：

>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         sorted, idx = x.sort()
>>>         ctx.mark_non_differentiable(idx)
>>>         ctx.save_for_backward(x, idx)
>>>         return sorted, idx
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):  # still need to accept g2
>>>         x, idx = ctx.saved_tensors
>>>         grad_input = torch.zeros_like(x)
>>>         grad_input.index_add_(0, idx, g1)
>>>         return grad_input

save_for_backward(*tensors)

保存給定的 Tensor，以便將來調用 backward()。

save_for_backward 最多應調用一次，無論是在 setup_context() 還是 forward() 方法中，並且僅能使用 Tensor。

所有打算在反向傳播中使用的 Tensor 都應使用 save_for_backward 保存（而不是直接在 ctx 上），以防止不正確的梯度和記憶體洩漏，並啟用已保存 Tensor 鉤子的應用。請參閱 torch.autograd.graph.saved_tensors_hooks。

請注意，如果中介 Tensor（既不是 forward() 的輸入，也不是輸出）被保存以供反向傳播使用，則您的自定義 Function 可能不支持雙重反向傳播。不支持雙重反向傳播的自定義 Function 應使用 @once_differentiable 修飾其 backward() 方法，以便執行雙重反向傳播會引發錯誤。如果您想支持雙重反向傳播，您可以根據反向傳播期間的輸入重新計算中介值，或將中介值作為自定義 Function 的輸出返回。有關更多詳細信息，請參閱雙重反向傳播教學。

在 backward() 中，可以透過 saved_tensors 屬性存取已儲存的 Tensor。在將它們返回給使用者之前，會進行檢查以確保它們沒有用於修改其內容的任何原地 (in-place) 操作。

參數也可以是 None。這是一個空操作 (no-op)。

有關如何使用此方法的更多詳細信息，請參閱 擴展 torch.autograd。

範例：

>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int):
>>>         w = x * z
>>>         out = x * y + y * z + w * y
>>>         ctx.save_for_backward(x, y, w, out)
>>>         ctx.z = z  # z is not a tensor
>>>         return out
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, grad_out):
>>>         x, y, w, out = ctx.saved_tensors
>>>         z = ctx.z
>>>         gx = grad_out * (y + y * z)
>>>         gy = grad_out * (x + z + w)
>>>         gz = None
>>>         return gx, gy, gz
>>>
>>> a = torch.tensor(1., requires_grad=True, dtype=torch.double)
>>> b = torch.tensor(2., requires_grad=True, dtype=torch.double)
>>> c = 4
>>> d = Func.apply(a, b, c)

save_for_forward(*tensors)

保存給定的 Tensor，以便將來調用 jvp()。

save_for_forward 最多應調用一次，無論是在 setup_context() 還是 forward() 方法中，並且所有參數都應該是 Tensor。

在 jvp() 中，可以透過 saved_tensors 屬性存取已儲存的物件。

參數也可以是 None。這是一個空操作 (no-op)。

有關如何使用此方法的更多詳細信息，請參閱 擴展 torch.autograd。

範例：

>>> class Func(torch.autograd.Function):
>>>     @staticmethod
>>>     def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int):
>>>         ctx.save_for_backward(x, y)
>>>         ctx.save_for_forward(x, y)
>>>         ctx.z = z
>>>         return x * y * z
>>>
>>>     @staticmethod
>>>     def jvp(ctx, x_t, y_t, _):
>>>         x, y = ctx.saved_tensors
>>>         z = ctx.z
>>>         return z * (y * x_t + x * y_t)
>>>
>>>     @staticmethod
>>>     def vjp(ctx, grad_out):
>>>         x, y = ctx.saved_tensors
>>>         z = ctx.z
>>>         return z * grad_out * y, z * grad_out * x, None
>>>
>>>     a = torch.tensor(1., requires_grad=True, dtype=torch.double)
>>>     t = torch.tensor(1., dtype=torch.double)
>>>     b = torch.tensor(2., requires_grad=True, dtype=torch.double)
>>>     c = 4
>>>
>>>     with fwAD.dual_level():
>>>         a_dual = fwAD.make_dual(a, t)
>>>         d = Func.apply(a_dual, b, c)

set_materialize_grads(value)

設定是否實體化梯度 Tensor。預設值為 True。

這應該僅從 setup_context() 或 forward() 方法中調用。

如果 True，則未定義的梯度 Tensor 將在調用 backward() 和 jvp() 方法之前擴展為充滿零的 Tensor。

範例：

>>> class SimpleFunc(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         return g1 + g2  # No check for None necessary
>>>
>>> # We modify SimpleFunc to handle non-materialized grad outputs
>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         ctx.set_materialize_grads(False)
>>>         ctx.save_for_backward(x)
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         x, = ctx.saved_tensors
>>>         grad_input = torch.zeros_like(x)
>>>         if g1 is not None:  # We must check for None now
>>>             grad_input += g1
>>>         if g2 is not None:
>>>             grad_input += g2
>>>         return grad_input
>>>
>>> a = torch.tensor(1., requires_grad=True)
>>> b, _ = Func.apply(a)  # induces g2 to be undefined

static setup_context(ctx, inputs, output)

有兩種定義 autograd.Function 的前向傳遞 (forward pass) 的方法。

任一種方式：

1. 使用簽章 forward(ctx, *args, **kwargs) 覆寫 forward。 setup_context 不會被覆寫。為 backward 設定 ctx 發生在 forward 內部。

2. 使用簽章 forward(*args, **kwargs) 覆寫 forward，並覆寫 setup_context。為 backward 設定 ctx 發生在 setup_context 內部（而不是在 forward 內部）。

詳情請參閱 torch.autograd.Function.forward() 和 擴展 torch.autograd。

傳回型別

Any

static vjp(ctx, *grad_outputs)

定義一個使用反向模式自動微分來微分運算的公式。

這個函數需要被所有子類別覆寫。（定義這個函數等同於定義 vjp 函數。）

它必須接受一個 context ctx 作為第一個引數，接著是和 forward() 傳回的一樣多的輸出（對於 forward 函數的非張量輸出，將傳入 None），並且它應該傳回和 forward() 的輸入一樣多的張量。每個引數是相對於給定輸出的梯度，並且每個傳回的值應該是相對於相應輸入的梯度。如果輸入不是 Tensor，或者是不需要梯度的 Tensor，您可以直接將 None 作為該輸入的梯度傳入。

context 可用於檢索在 forward 傳遞期間儲存的張量。它還有一個屬性 ctx.needs_input_grad，它是一個布林值元組，表示每個輸入是否需要梯度。例如，如果 forward() 的第一個輸入需要計算相對於輸出的梯度，則 backward() 將具有 ctx.needs_input_grad[0] = True。

傳回型別

Any

static vmap(info, in_dims, *args)

定義在 torch.vmap() 下的此 autograd.Function 的行為。

為了讓 torch.autograd.Function() 支援 torch.vmap()，您必須覆寫此靜態方法，或者將 generate_vmap_rule 設定為 True (您不能同時做這兩件事)。

如果您選擇覆寫此 staticmethod：它必須接受

• 一個 info 物件作為第一個參數。info.batch_size 指定了正在 vmap 的維度的大小，而 info.randomness 是傳遞給 torch.vmap() 的隨機性選項。

• 一個 in_dims tuple 作為第二個參數。對於 args 中的每個 arg，in_dims 都有一個對應的 Optional[int]。如果 arg 不是 Tensor 或 arg 沒有被 vmap，則它是 None，否則它是一個整數，指定 Tensor 的哪個維度正在被 vmap。

• *args，它與 forward() 的 args 相同。

vmap staticmethod 的回傳值是一個 (output, out_dims) 的 tuple。類似於 in_dims，out_dims 應該與 output 具有相同的結構，並且每個輸出包含一個 out_dim，用於指定輸出是否具有 vmap 維度以及它所在的索引。

更多詳情請參閱 使用 autograd.Function 擴展 torch.func。