InplaceFunction

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

此類僅用於向後相容。對於任何新用例，請使用 Function 代替此項。

static backward(ctx, *grad_outputs)

定義使用反向模式自動微分來計算操作導數的公式。

所有子類都應覆蓋此函式。（定義此函式等同於定義 vjp 函式。）

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

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

返回型別

任意

static forward(*args, **kwargs)

定義自定義 autograd 函式的前向傳播。

所有子類都應覆蓋此函式。定義 forward 有兩種方式

用法 1（合併 forward 和 ctx）

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

• 它必須接受一個上下文 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() 靜態方法來處理 ctx 物件的設定。output 是 forward 的輸出，inputs 是 forward 輸入的元組。

• 有關更多詳細資訊，請參閱 擴充套件 torch.autograd

上下文可用於儲存可在反向傳播期間檢索的任意資料。張量不應直接儲存在 ctx 上（儘管為了向後相容目前並未強制執行此規則）。相反，如果張量用於 backward（等同於 vjp），應使用 ctx.save_for_backward() 儲存；如果用於 jvp，應使用 ctx.save_for_forward() 儲存。

返回型別

任意

static jvp(ctx, *grad_inputs)

定義使用前向模式自動微分來計算操作導數的公式。

所有子類都應覆蓋此函式。它必須接受一個上下文 ctx 作為第一個引數，後跟與 forward() 輸入數量相同的引數（forward 函式的非張量輸入將傳入 None），並且它應該返回與 forward() 輸出數量相同的張量。每個引數是相對於給定輸入的梯度，每個返回值應該是相對於相應輸出的梯度。如果輸出不是張量或函式相對於該輸出不可微分，您可以直接傳入 None 作為該輸入的梯度。

您可以使用 ctx 物件將前向傳播中的任何值傳遞給此函式。

返回型別

任意

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() 方法中呼叫，所有引數都應該是張量輸出。

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

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

>>> 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)

儲存給定的張量，以便將來呼叫 backward()。

save_for_backward 最多隻能呼叫一次，可以在 setup_context() 或 forward() 方法中呼叫，且只能用於張量。

所有用於反向傳播的張量都應使用 save_for_backward 儲存（而不是直接儲存在 ctx 上），以防止梯度錯誤和記憶體洩漏，並啟用 saved tensor hooks 的應用。參閱 torch.autograd.graph.saved_tensors_hooks。

請注意，如果中間張量（即 forward() 的輸入和輸出之外的張量）被儲存用於反向傳播，您的自定義 Function 可能不支援二次反向傳播（double backward）。不支援二次反向傳播的自定義 Functions 應使用 @once_differentiable 裝飾器裝飾其 backward() 方法，以便執行二次反向傳播時引發錯誤。如果您想支援二次反向傳播，可以在 backward 期間基於輸入重新計算中間量，或將中間量作為自定義 Function 的輸出返回。有關更多詳細資訊，請參閱 二次反向傳播教程。

在 backward() 中，可以透過 saved_tensors 屬性訪問儲存的張量。在將它們返回給使用者之前，會進行檢查以確保它們未在任何修改其內容的就地操作中使用。

引數也可以是 None。這相當於空操作。

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

儲存給定的張量，以便將來呼叫 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)

set_materialize_grads(value)

設定是否具體化梯度張量。預設為 True。

此函式只能在 setup_context() 或 forward() 方法中呼叫。

如果為 True，則在呼叫 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 的前向傳播有兩種方式。

或者

1. 使用 forward(ctx, *args, **kwargs) 簽名覆蓋 forward。不覆蓋 setup_context。為反向傳播設定 ctx 在 forward 內部進行。

2. 使用 forward(*args, **kwargs) 簽名覆蓋 forward 並覆蓋 setup_context。為反向傳播設定 ctx 在 setup_context 內部進行（而不是在 forward 內部）

有關更多詳細資訊，請參閱 torch.autograd.Function.forward() 和 擴充套件 torch.autograd。

返回型別

任意

static vjp(ctx, *grad_outputs)

定義使用反向模式自動微分來計算操作導數的公式。

所有子類都應覆蓋此函式。（定義此函式等同於定義 vjp 函式。）

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

上下文可用於檢索前向傳播期間儲存的張量。它還有一個屬性 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() 的 randomness 選項。

• 一個 in_dims 元組作為第二個引數。對於 args 中的每個 arg，in_dims 都有一個對應的 Optional[int]`. 如果 arg 不是張量或沒有進行 vmap 操作，則為 `None`；否則，它是一個整數，指定正在對張量的哪個維度進行 vmap 操作。

• *args，與傳遞給 forward() 的引數相同。

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

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