关于pytorch requires_grad一二事

2024年3月12日 259次阅读 来源: 东方翼轩

这两天一直迷惑pytorch中关于requires_grad的设置，特别是在写cuda extension时。今天在pytorch github上看到个issue解释的很不错，姑且记录在这以备后用。

我的问题是在写cuda扩展时，我会在python autograd.Function的forward部分先声明好变量，就是先分配好内存，然后传入cuda的函数，cuda内做自定义操作，类似：

class myfunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        out = torch.zeros_like(x)
        cuda_foward(x, out)
        ctx.save_for_backward(x)
        return out
    @staticmethod
    def backward(ctx, d_out):
        x, = ctx.saved_tensors
        d_x = torch.zeros_like(x)
        cuda_backward(d_x, x, d_out)
        return d_x

OK, 那问题是forward内部声明的out要不要设定requires_grad的呢？

详见该issue：https://github.com/pytorch/pytorch/issues/17849

  1. 当声明 `out = torch.zeros_like(x)` 时，尽管x是requires_grad的，但out不是
  2. 通过autograd.Function的返回值，看来是会被设为 requires_grad的，除非在forward内设定ctx.mark_non_differentiable(out)

OK, 主要是关于pytorch的autograd机制目前还没时间细致了解，所以先把这些实用内容留在这儿，以备后需。

    原文作者：东方翼轩
    原文地址: https://zhuanlan.zhihu.com/p/81591380
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞