評價此頁

torch.autograd.function.FunctionCtx.mark_non_differentiable#

FunctionCtx.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