torch.autograd.Function的使用
（个人理解仅供参考）
1 什么情况下使用
自己定义的网络结构，没有现成的，就得手写forward和backward
2 怎么使用
2.1 forward
前向传播的表达式
2.2 backward
求导结果
2.3 举例
前向传播表达式：y = w * x + b
假设f()是我们关于y的loss函数，那么z = f(y)即为loss值
现在要求loss对w、x、b的偏导（假设只有一层）：
dz/dx = dz/dy * dy/dx = dz/dy * w
dz/dw = dz/dy * dy/dw = dz/dy * x
dz/db = dz/dy * dy/db = dz/dy * 1
好在dz/dy不用我们再求了，它就是 backward 的参数grad_output。那么grad_output是从哪来的呢？其实就是 forward 会 return output 给 backward ，至于 backward 怎么把 output 变为 grad_output 就不用细究了。
所以：
dz/dx = grad_output * w
dz/dw = grad_output * x
dz/db = grad_output * 1
因此，对于y = w * x + b，我们的代码为：
import torch
from torch.autograd import Function
class MultiplyAdd(Function):
@staticmethod
def forward(ctx, w, x, b):
ctx.save_for_backward(w, x) # 保存参数
output = w * x + b
return output # 传给backward
@staticmethod
def backward(ctx, grad_output):
w, x = ctx.saved_tensors
grad_w = grad_output * x
grad_x = grad_output * w
grad_b = grad_output * 1
return grad_w, grad_x, grad_b # 传给forward
Linear = MultiplyAdd.apply
2.4 模板
"""
# 使用autograd.Function进行扩展的一个模板
class My_Function(Function):
def forward(self, inputs, parameters):
self.saved_for_backward = [inputs, parameters]
# output = [对输入和参数进行的操作，其实就是前向运算的函数表达式]
return output
def backward(self, grad_output):
inputs, parameters = self.saved_tensors # 或self.saved_variables
# grad_input = [求函数forward(input)关于 parameters 的导数，其实就是反向运算的导数表达式] * grad_output
return grad_input
"""
2.5 验证
验证的话需要使用torch.autograd.gradcheck，给上我的完整代码，验证部分在最后：
import torch
from torch.autograd import Function, gradcheck
"""
# 使用autograd.Function进行扩展的一个模板
class My_Function(Function):
def forward(self, inputs, parameters):
self.saved_for_backward = [inputs, parameters]
# output = [对输入和参数进行的操作，其实就是前向运算的函数表达式]
return output
def backward(self, grad_output):
inputs, parameters = self.saved_tensors # 或者是self.saved_variables
# grad_input = [求函数forward(input)关于 parameters 的导数，其实就是反向运算的导数表达式] * grad_output
return grad_input
"""
class MultiplyAdd(Function):
@staticmethod
def forward(ctx, w, x, b):
ctx.save_for_backward(w, x)
output = w * x + b
return output
@staticmethod
def backward(ctx, grad_output):
w, x = ctx.saved_tensors
grad_w = grad_output * x
grad_x = grad_output * w
grad_b = grad_output * 1
return grad_w, grad_x, grad_b
Linear = MultiplyAdd.apply
x = torch.ones(1, requires_grad=True, dtype=torch.float64)
w = torch.rand(1, requires_grad=True, dtype=torch.float64)
b = torch.rand(1, requires_grad=True, dtype=torch.float64)
# print("start forward...")
# z = MultiplyAdd.apply(w, x, b)
# print("start backward...")
# z.backward()
#
# print(x.grad, w.grad, b.grad)
test = gradcheck(Linear, (x, w, b), eps=1e-6)
print(test)
3 存疑
现在我只是会用了这个，但是如果是两层的全连接层，这段代码是怎么工作的？这个问题我还没想明白，留个坑