快捷方式

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源