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