pytorch自定义函数实现自动梯度

Motivation

构建模型有时需要使用自定义的函数，为了不影响模型的反向传播，需要实现自动梯度计算（即把自定义函数嵌入计算图）。

---

实现

要点：

1. 将函数定义为类，需继承自torch.autograd.Function类
2. 需实现两个静态方法：forward()和backward()，分别对应前向传播和反向传播
3. 函数使用前需调用apply方法从而嵌入计算图，实现自动求导

用一个例子来说明:
假设我们要实现一个多项式拟合模型： $y=a+b{P}_2(cx+d)$ ，用来拟合正切函数，其中$P_2$为一个二次函数：$P_2(x)=x^2+2x$，需要我们自定义实现

import torch
import math

dtype = torch.float
#构建训练集,使用随机数据
x = torch.linspace(-math.pi, math.pi, 2000, dtype = dtype)
y = torch.tan(x)

#将权重值a,b,c,d随机初始化，设置requires_grad=True以实现自动求偏导
a = torch.randn((), dtype = dtype, requires_grad = True)
b = torch.randn((), dtype = dtype, requires_grad = True)
c = torch.randn((), dtype = dtype, requires_grad = True)
d = torch.randn((), dtype = dtype, requires_grad = True)

#实现P2
class P2(torch.autograd.Function): 
	#定义forward方法，以在前向传播中实现函数的功能
	@staticmethod
	def forward(ctx, input):
		#保存input值，反向传播计算梯度时会用到
		ctx.save_for_backward(input)
		return input**2 + 2*input
	
	#定义backward方法
	@staticmethod
	def backward(ctx,grad_output): 
		#输入参数grad_output是目标函数对输出结果的梯度
		#这里需要计算函数输出值相对于输入input的梯度u，
		#返回grad_output与u的乘积以实现链式法则
		input, = ctx.saved_tensors
		return grad_output * (2 * inptu + 2)

learning_rate = 1e-6
#全量学习
for t in range(2000):
	#对刚刚定义的函数fun调用apply方法，使其嵌入计算图
	p2 = P2.apply
	#前向传播
	y_pred = a + b * p2(c * x + d)
	#定义损失函数
	loss = (y_pred - y).pow(2).sum()
	if t%100 == 99:
		print(t, loss.item())
	#反向传播计算梯度
	loss.backward()
	#更新参数
	#这里每个参数的梯度自定在反向传播时自动计算好并保存在param.grad中了
	with torch.no_grad():
		a -= learning_rate * a.grad
		b -= learning_rate * b.grad
		c -= learning_rate * c.grad
		d -= learning_rate * d.grad

		#将梯度重置为0
		a.grad = b.grad = c.grad = d.grad = None
print(f'Result: y = {a.item()} + {b.item()} * P2({c.item} x + {d.item()}')
	

运行结果：

（由于参数选择很随意，且权重是随机初始化的，所以拟合效果不好）

---

参考资料

链接：Learning pytorch with examples