pytorch的自定义拓展之（三）——torch.autograd.Function的简单定义与案例

前言：详细过程可以参考前面的两篇文章：

pytorch的自定义拓展之（二）——torch.autograd.Function完成自定义层

pytorch的自定义拓展之（一）——torch.nn.Module和torch.autograd.Function

 

一、使用autograd.Function 定义新的自动求导函数

 

在底层，每一个原始的自动求导运算实际上是两个在Tensor上运行的函数。其中，forward函数计算从输入Tensors获得的输出Tensors。而backward函数接收输出Tensors对于某个标量值的梯度，并且计算输入Tensors相对于该相同标量值的梯度。

在PyTorch中，我们可以很容易地通过定义torch.autograd.Function的子类并实现forward和backward函数，来定义自己的自动求导运算。之后我们就可以使用这个新的自动梯度运算符了。然后，我们可以通过构造一个实例并像调用函数一样，传入包含输入数据的tensor调用它，这样来使用新的自动求导运算。

这个例子中，我们自定义一个自动求导函数来展示ReLU的非线性。并用它实现我们的两层网络：

import torch

class MyReLU(torch.autograd.Function):
"""
我们可以通过建立torch.autograd的子类来实现我们自定义的autograd函数，
并完成张量的正向和反向传播。
"""
@staticmethod
def forward(ctx, x):
"""
在正向传播中，我们接收到一个上下文对象和一个包含输入的张量；
我们必须返回一个包含输出的张量，
并且我们可以使用上下文对象来缓存对象，以便在反向传播中使用。
"""
ctx.save_for_backward(x)
return x.clamp(min=0)

@staticmethod
def backward(ctx, grad_output):
"""
在反向传播中，我们接收到上下文对象和一个张量，
其包含了相对于正向传播过程中产生的输出的损失的梯度。
我们可以从上下文对象中检索缓存的数据，
并且必须计算并返回与正向传播的输入相关的损失的梯度。
"""
x, = ctx.saved_tensors
grad_x = grad_output.clone()
grad_x[x < 0] = 0
return grad_x


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# N是批大小；D_in 是输入维度；
# H 是隐藏层维度；D_out 是输出维度
N, D_in, H, D_out = 64, 1000, 100, 10

# 产生输入和输出的随机张量
x = torch.randn(N, D_in, device=device)
y = torch.randn(N, D_out, device=device)

# 产生随机权重的张量
w1 = torch.randn(D_in, H, device=device, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, requires_grad=True)

learning_rate = 1e-6
for t in range(500):
# 正向传播：使用张量上的操作来计算输出值y；
# 我们通过调用 MyReLU.apply 函数来使用自定义的ReLU
y_pred = MyReLU.apply(x.mm(w1)).mm(w2)

# 计算并输出loss
loss = (y_pred - y).pow(2).sum()
print(t, loss.item())

# 使用autograd计算反向传播过程。
loss.backward()

with torch.no_grad():
# 用梯度下降更新权重
w1 -= learning_rate * w1.grad
w2 -= learning_rate * w2.grad

# 在反向传播之后手动清零梯度
w1.grad.zero_()
w2.grad.zero_()