用张量广播机制实现神经网络反向传播

正向传播

要想了解反向传播，先要了解正向传播：正向传播的每一步是，用一个或不少输入生成一个输出。

反向传播

反向传播的做用是计算模型参数的偏导数。再具体一点，反向传播的每个step就是：已知正向传播的输入自己，和输出的偏导数，求出每一个输入的偏导数的过程。

反向传播既简单，又复杂：

• 它的原理很简单：链式法则求偏导。
• 它的公式又很复杂：由于它的公式看起来真的很复杂。

模型的参数

反向传播就是计算模型的参数的偏导数，因此介绍一下模型的参数：

• 模型里有不少参数，参数的本质是张量，能够把张量当作多维数组，也能够把张量当作一颗树。
• 张量有形状，张量的偏导数是一个一样形状的张量。

线性函数的反向传播

线性函数就是 y = wx + b，咱们输入x，w，和 b 就能获得y。y是咱们算出来的，这个算y的过程就是正向传播。

咱们规定字母后面加 .g 表示偏导数，如 y.g 就是y的偏导数，w.g 就是w的偏导数。

那么咱们的目的，就是根据 x, w, b 和 y.g 的值，分别算出 w,x,和b的偏导数，而这个过程，就是反向传播。

为了便于说明，咱们假设了每一个变量的形状: x(1000, 784), w(784, 50), b(50), y(1000, 50)。

计算 x.g

y = wx + b 对 x 求偏导 得 w，即咱们要用 w 和 y.g 计算出 x.g。

w 的形状是 (784, 50)，y.g的形状跟y相同，是(1000, 50)，如何用这两个形状凑出 x.g 的(1000, 784)？

emmm，很简单，就是这样，而后那样，就好了。看玩笑的。。其实就是 y.g 中间加一维，变成 (1000, 1, 50) ，而后再跟 w 搞一下，获得一个 (1000, 784, 50) 的形状，再把最后一维消去，就获得 (1000, 784) 的形状了。

即:
x.g = (y.g.unsqueeze(1) * w).sum(dim=-1)

计算 w.g

同理咯，y = wx + b 对 w 求偏导 得 x，即咱们要用 x 和 y.g 计算出 w.g。

x 的形状是 (1000, 784)，y.g的形状跟y相同，是(1000, 50)，如何用这两个形状凑出 w.g` 的(784, 50)？

先将 x 最后加一维，变成 (1000, 784, 1)，再将 y.g 中间加一维，变成 (1000, 1, 50)，这俩搞一下，变成 (1000, 784, 50)，再把开头的那一维消去，就变成 (784, 50)了。

即:
w.g = (x.unsqueeze(-1) * y.g.unsqueeze(1)).sum(dim=0)

计算 b.g

y = wx + b 对 b 求偏导 得常数 1，因此直接用形状为(1000, 50)的y.g来凑出形状为(50)的b.g就能够了。

那么就很是简单了，直接把(1000, 50)消去最开始的那一维就能获得(50)，即：

b.g = y.g.sum(0)

线性函数的反向传播代码

已知线性函数的输入是 inp，输出是 out，计算过程用到的两个参数是 w和b，则反向传播的代码以下：

def back_lin(inp, w, b, out):
    inp.g = (out.g.unsqueeze(1) * w).sum(dim=-1)
    w.g = (inp.unsqueeze(-1) * out.g.unsqueeze(1)).sum(dim=0)
    b.g = out.g.sum(0)

relu函数的反向传播

relu函数表示起来很简单，就是 max(x, 0)，可是在 pytorch 中这样写是行不通的，因此用这面这个函数表示：

def relu(x):
    return x.clamp_min(0)

其反向传播表示为:

def back_relu(inp, out):
    return (inp > 0).float() * out.g

mse函数的反向传播

mse函数用代码表示为:

def mse(pred, target):
    return (pred.squeeze(dim=-1)-target).pow(2).mean()

其反向传播则是:

def back_mse(pred, target):
    return 2. * (pred.squeeze(dim=-1) - target).unsqueeze(dim=-1) / pred.shape[0]

测试

假设咱们的模型结果为：输入一个x，进行一次线性变换，再通过一次relu，而后再通过一次线性变换获得结果。

先随机生成 输入、输出和各个参数:

# 伪造输入和答案
import torch
torch.manual_seed(0)
input_ = torch.randn(1000, 784).requires_grad_(True)  # 输入
target = torch.randn(1000)  # 答案
# 建立其它参数
w1 = torch.randn(784, 50).requires_grad_(True)
b1 = torch.randn(50).requires_grad_(True)
w2 = torch.randn(50, 1).requires_grad_(True)
b2 = torch.randn(1).requires_grad_(True)

正向传播获得模型的输出:

l1 = input_ @ w1 + b1
l2 = relu(l1)
output = l2 @ w2 + b2
loss = mse(output, target)

反向传播:

back_mse(output, target)
back_lin(l2, w2, b2, output)
back_relu(l1, l2)
back_lin(input_, w1, b1, l1)

此时 w1.g，b1.g和 w2.g，b2.g均已求出。

而后用pytorch自带的反向传播求一下梯度:

# 先保存一下手动求的梯度
w1g = w1.g.clone()
b1g = b1.g.clone()
w2g = w2.g.clone()
b2g = b2.g.clone()

input_ = input_.clone().requires_grad_(True)
w1 = w1.clone().requires_grad_(True)
b1 = b1.clone().requires_grad_(True)
w2 = w2.clone().requires_grad_(True)
b2 = b2.clone().requires_grad_(True)

l1 = input_ @ w1 + b1
l2 = relu(l1)
output = l2 @ w2 + b2
loss = mse(output, target)

loss.backward()

此时对比一下咱们手动求得的梯度和调用系统函数求得的梯度，发现两者是相等的：

def is_same(a, b):
    return (a - b).max() < 1e-4

is_same(w1g, w1.grad), is_same(b2g, b2.grad), is_same(w2g, w2.grad), is_same(b2g, b2.grad)
"""输出
(tensor(True), tensor(True), tensor(True), tensor(True))
"""

总结

借助简单的求导和张量的广播机制，就能够推导实现神经网络的反向传播。