NestedIOFunction

class torch.autograd.function.NestedIOFunction(*args, **kwargs)

此類僅出於向後相容性原因而存在。在任何新用例中，請使用 Function 而不是此類。

backward(*gradients)

共享的 backward 工具。

返回型別

任何

backward_extended(*grad_output)

使用者定義的 backward。

forward(*args)

共享的 forward 工具。

返回型別

任何

forward_extended(*input)

使用者定義的 forward。

static jvp(ctx, *grad_inputs)

定義使用前向模式自動微分來區分操作的公式。

此函式應由所有子類重寫。它必須接受一個上下文 ctx 作為第一個引數，然後是 forward() 接收到的輸入數量（對於 forward 函式的非張量輸入將傳入 None），並且它應該返回與 forward() 的輸出數量相同的張量。每個引數是關於給定輸入的梯度，每個返回值應該是關於相應輸出的梯度。如果輸出不是張量或函式對該輸出不可微，則只需為該輸入傳遞 None 作為梯度。

You can use the ctx object to pass any value from the forward to this functions.

返回型別

任何

mark_dirty(*args, **kwargs)

參見 Function.mark_dirty()。

mark_non_differentiable(*args, **kwargs)

參見 Function.mark_non_differentiable()。

save_for_backward(*args)

參見 Function.save_for_backward()。

save_for_forward(*tensors)

Save given tensors for a future call to jvp().

save_for_forward 最多應呼叫一次，在 setup_context() 或 forward() 方法中，並且所有引數都應該是張量。

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

引數也可以是 None。這不會執行任何操作。

有關如何使用此方法的更多詳細資訊，請參閱 擴充套件 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)

property saved_tensors

參見 Function.saved_tensors()。

set_materialize_grads(value)

Set whether to materialize grad tensors. Default is True.

這應僅從 setup_context() 或 forward() 方法中呼叫。

如果為 True，未定義的 grad 張量將在呼叫 backward() 和 jvp() 方法之前擴充套件為全零張量。

示例

>>> 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 的前向傳遞。

Either

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 以獲取更多詳細資訊。

返回型別

任何

static vjp(ctx, *grad_outputs)

定義使用反向模式自動微分來區分操作的公式。

此函式應被所有子類重寫。（定義此函式等同於定義 vjp 函式。）

它必須接受一個上下文 ctx 作為第一個引數，然後是 forward() 返回的輸出數量（對於 forward 函式的非張量輸出將傳入 None），並且它應該返回與 forward() 的輸入數量相同的張量。每個引數是關於給定輸出的梯度，每個返回值應該是關於相應輸入的梯度。如果輸入不是張量或是一個不需要梯度的張量，則可以為該輸入傳遞 None 作為梯度。

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

返回型別

任何

static vmap(info, in_dims, *args)

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

要使 torch.autograd.Function() 支援 torch.vmap()，您必須覆蓋此靜態方法，或者將 generate_vmap_rule 設定為 True（您不能同時執行這兩項）。

如果您選擇重寫此靜態方法：它必須接受

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

• 第二個引數是一個 in_dims 元組。對於 args 中的每個 arg，in_dims 有一個相應的 Optional[int]。如果 arg 不是 Tensor 或 arg 不被 vmap，則為 None，否則，它是一個指定 Tensor 的哪個維度被 vmap 的整數。

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

vmap 靜態方法的返回值是一個元組 (output, out_dims)。與 in_dims 類似，out_dims 的結構應與 output 相同，並且每個輸出都包含一個 out_dim，指定輸出是否具有 vmap 的維度以及在該維度中的索引。

有關更多詳細資訊，請參閱 使用 autograd.Function 擴充套件 torch.func。