Python PyTorch FunctionCtx.mark_dirty用法及代碼示例

本文簡要介紹python語言中 torch.autograd.function.FunctionCtx.mark_dirty 的用法。

用法:

FunctionCtx.mark_dirty(*args)

將給定張量標記為在就地操作中修改。

這應該最多調用一次，隻能從內部調用 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