玩轉Pytorch的Function類

---

前言

 pytorch提供了autograd自動求導機制，而autograd實現自動求導實質上通過Function類實現的。而習慣搭積木的夥伴平時也不寫backward。造成需要拓展算子情况便會手足無措。本文從簡單例子入手，學習實現一個Function類最基本的要素，同時還會涉及一些注意事項，最後在結合一個實戰來學習Function類的使用。

1、y=w*x+b

import torch
from torch.autograd import Function

# y = w*x + b 的一個前向傳播和反向求導
class Mul(Function):
    @staticmethod
    def forward(ctx, w, x, b, x_requires_grad = True): # ctx可以理解為元祖，用來存儲梯度的中間緩存變量。
        ctx.save_for_backward(w,b)     # 因為dy/dx = w; dy/dw = x ; dy/db = 1；為了後續反向傳播需要保存中間變量w,x
        output = w*x + b
        return output
    @staticmethod
    def backward(ctx,grad_outputs):    # 此處grad_outputs 具體問題具體分析
        w = ctx.saved_tensors[0]      # 取出ctx中保存的 w = 2
        b = ctx.saved_tensors[1]      # 取出ctx中保存的 b = 3
        grad_w = grad_outputs * x     # 1 * 1 = 1
        grad_x = grad_outputs * w     # 1 * 2 = 2
        grad_b = grad_outputs * 1     # 1 * 1 = 1
        return grad_w, grad_x, grad_b, None  # 返回的參數和forward的參數一一對應，對於參數x_requires_grad不必求梯度則直接返回None。

if __name__ == '__main__':
    x = torch.tensor(1.,requires_grad=True)
    w = torch.tensor(2.,requires_grad=True)
    b = torch.tensor(3., requires_grad=True)
    y = Mul.apply(w,x,b)              # y = w*x + b = 2*1 + 3 = 5
    print('forward:', y)
    # 寫法一
    loss = y.sum()                    # 轉成標量
    loss.backward()                   # 反向傳播：因為 loss = sum(y),故grad_outputs = dloss/dy = 1,可以省略不寫
    print('寫法一的梯度：',x.grad, w.grad, b.grad)      # tensor(2.) tensor(1.) tensor(1.)

 這裏簡單說下：代碼中注釋有問題歡迎留言評論。其中y=w*x+b。前向傳播容易理解。這裏令人困惑的應該是ctx這個東西，其實可以將其理解為一個元祖，通過方法save_for_backward()來保存前向傳播的中間緩存變量，為後續反向傳播提供條件。而在反向傳播中，首先從ctx中通過調用方法saved_tensors[]來得到w,b。之後各個參數的梯度：dy/dx = w; dy/dw = x; dy/db = 1。
 另外，在反向傳播中，令人困惑就是參數grad_outputs。其實這個參數的值跟類調用完之後有關。在代碼中，使用loss.backward()，可以看見傳入的參數是個空。這是因為在計算完前向傳播得到y之後，loss = y.sum()，即grad_outputs = dloss/dy = 1; 而在torch中，可以省略不寫。故此處的grad_outputs=1.
 當然，我們也可以明示的傳參進去。

    # 寫法一
    loss1 = y.sum()                    # 轉成標量
    loss1.backward()                   # 反向傳播：因為 loss = sum(y),故grad_outputs = dloss/dy = 1,可以省略不寫
    print('寫法一的梯度：',x.grad, w.grad, b.grad)      # tensor(2.) tensor(1.) tensor(1.)
    # 寫法二
    loss2 = y.sum()
    loss2.backward(torch.tensor(1.))
    print('寫法二的梯度：',x.grad, w.grad, b.grad)      # tensor(4.) tensor(2.) tensor(2.)

 但是此時報錯了，報錯信息如下：

RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.

 大體意思說同一個計算圖不能反向傳播兩次。因為，在調用第一次backward之後，計算圖就銷毀了。所以需要通過設置參數retain_graph參數保存計算圖，更改後代碼如下：

    # 寫法一
    loss1 = y.sum()                    # 轉成標量
    loss1.backward(retain_graph = True)                   # 反向傳播：因為 loss = sum(y),故grad_outputs = dloss/dy = 1,可以省略不寫
    print('寫法一的梯度：',x.grad, w.grad, b.grad)      # tensor(2.) tensor(1.) tensor(1.)
    # 寫法二
    loss2 = y.sum()
    loss2.backward(torch.tensor(1.))
    print('寫法二的梯度：',x.grad, w.grad, b.grad)      # tensor(4.) tensor(2.) tensor(2.)

 不幸的是，此時寫法二和寫法一的梯度計算結果不一致，發現寫法二的梯度是寫法一梯度的兩倍。是因為在pytorch中兩次不同loss在反傳梯度時在葉子節點梯度是累加的。因此，我們在損失二傳播之間需要將損失一的梯度清0。代碼如下：

    # 寫法一
    loss1 = y.sum()                    # 轉成標量
    loss1.backward(retain_graph = True)                   # 反向傳播：因為 loss = sum(y),故grad_outputs = dloss/dy = 1,可以省略不寫
    print('寫法一的梯度：',x.grad, w.grad, b.grad)      # tensor(2.) tensor(1.) tensor(1.)
    # 葉子節點梯度清0
    x.grad.zero_()
    w.grad.zero_()
    b.grad.zero_()
    # 寫法二
    loss2 = y.sum()
    loss2.backward(torch.tensor(1.))
    print('寫法二的梯度：',x.grad, w.grad, b.grad)      # tensor(2.) tensor(1.) tensor(1.)

 OK，大功告成。完整代碼如下：

import torch
from torch.autograd import Function

# y = w*x + b 的一個前向傳播和反向求導
class Mul(Function):
    @staticmethod
    def forward(ctx, w, x, b, x_requires_grad = True): # ctx可以理解為元祖，用來存儲梯度的中間緩存變量。
        ctx.save_for_backward(w,b)     # 因為dy/dx = w; dy/dw = x ; dy/db = 1；為了後續反向傳播需要保存中間變量w,x
        output = w*x + b
        return output
    @staticmethod
    def backward(ctx,grad_outputs):    # 此處grad_outputs 具體問題具體分析
        w = ctx.saved_tensors[0]      # 取出ctx中保存的 w = 2
        b = ctx.saved_tensors[1]      # 取出ctx中保存的 b = 3
        grad_w = grad_outputs * x     # 1 * 1 = 1
        grad_x = grad_outputs * w     # 1 * 2 = 2
        grad_b = grad_outputs * 1     # 1 * 1 = 1
        return grad_w, grad_x, grad_b, None  # 返回的參數和forward的參數一一對應，對於參數x_requires_grad不必求梯度則直接返回None。

if __name__ == '__main__':
    x = torch.tensor(1.,requires_grad=True)
    w = torch.tensor(2.,requires_grad=True)
    b = torch.tensor(3., requires_grad=True)
    y = Mul.apply(w,x,b)              # y = w*x + b = 2*1 + 3 = 5
    print('forward:', y)
    # 寫法一
    loss1 = y.sum()                    # 轉成標量
    loss1.backward(retain_graph = True)                   # 反向傳播：因為 loss = sum(y),故grad_outputs = dloss/dy = 1,可以省略不寫
    print('寫法一的梯度：',x.grad, w.grad, b.grad)      # tensor(2.) tensor(1.) tensor(1.)
    # 葉子節點梯度清0
    x.grad.zero_()
    w.grad.zero_()
    b.grad.zero_()
    # 寫法二
    loss2 = y.sum()
    loss2.backward(torch.tensor(1.))
    print('寫法二的梯度：',x.grad, w.grad, b.grad)      # tensor(4.) tensor(2.) tensor(2.)

2、進階：y=exp(x)*2

import torch
from torch.autograd import Function

class Exp(Function):
    @staticmethod
    def forward(ctx,x):
        output = x.exp()
        ctx.save_for_backward(output) # dy/dx = exp(x)
        return output
    @staticmethod
    def backward(ctx, grad_outputs):  # dloss/dx = grad_outputs* exp(x)
        output = ctx.saved_tensors[0]
        return output*grad_outputs

if __name__ == '__main__':
    x = torch.tensor(1.,requires_grad=True)
    y = Exp.apply(x)
    print(y)
    y = y * 2
    loss = y.sum()
    loss.backward()
    print(x.grad)

 唯一需要注意就是：dloss/dy = 1 * 2 = 2;因為loss = sum(2y)。

3、實戰：GuideReLU函數

  ReLU函數：y=max(x,0)，反傳梯度時僅x>0的比特置才有梯度，且梯度值為1.因為y=x,所以dy/dx=1；而GuideReLU是在ReLU基礎上，不僅x>0比特置才能反傳梯度，還要滿足梯度>0比特置才能反傳梯度。dloss/dx = dloss/dy * (x>0) * (grad_output>0)。代碼如下：

import torch
from torch.autograd import Function

class GuideReLU(Function):
    @staticmethod
    def forward(ctx,input):
        output = torch.clamp(input,min=0)
        ctx.save_for_backward(output)
        return output
    @staticmethod
    def backward(ctx, grad_outputs):  # dloss/dx = dloss/dy * !(x > 0) * (dloss/dy > 0)
        output = ctx.saved_tensors[0]  # dloss/dy
        return grad_outputs * (output>0).float()* (grad_outputs>0).float()

if __name__ == '__main__':
    x = torch.randn(2,3,requires_grad=True)
    print('input:',x)
    y = GuideReLU.apply(x)
    print('forward:',y)
    grad_y = torch.randn(2,3)
    y.backward(grad_y)               # 此處接收一個梯度數值，即grad_outputs
    print('grad_y:',grad_y)          # 即只有當輸入x和返回梯度grad_y同時>0比特置才有梯度值。
    print('grad_x:',x.grad)

4、梯度檢查：torch.autograd.gradcheck()

 pytorch提供了一個梯度檢查api，可以很方便檢測自己寫的傳播是否正確。

import torch
from torch.autograd import Function
class Sigmoid(Function):

    @staticmethod
    def forward(ctx, x):
        output = 1 / (1 + torch.exp(-x))
        ctx.save_for_backward(output)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        output, = ctx.saved_tensors
        grad_x = output * (1 - output) * grad_output
        return grad_x


test_input = torch.randn(4, requires_grad=True)  # tensor([-0.4646, -0.4403, 1.2525, -0.5953], requires_grad=True)
print(torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-3))

總結

 本篇是介紹pytorch反向傳導的第一篇，後續會介紹拓展C++/CUDA算子。若有問題歡迎+vx：wulele2541612007，拉你進群探討交流。