代码实现nadaraya-waston核回归 #51CTO博主之星评选#

Posted 2022-03-18 LolitaAnn

import torch
from torch import nn
from d2l import torch as d2l

最基础的导包，看不懂的python需要回炉重造，不建议继续往下看文章。

n_train = 50  # 训练样本数
x_train, _ = torch.sort(torch.rand(n_train) * 5)   # 训练样本的输入

def f(x):
    return 2 * torch.sin(x) + x**0.8

y_train = f(x_train) + torch.normal(0.0, 0.5, (n_train,))  # 训练样本的输出

x_test = torch.arange(0, 5, 0.1)  # 测试样本
y_truth = f(x_test)  # 测试样本的真实输出
n_test = len(x_test)  # 测试样本数

生成随机数据作为数据集。 设置要生成的训练样本数量为50个。

• torch.rand(n_train) * 5：生成$[0, 5)$之间的50个数据x。

• 使用torch.sort对其进行排序。这里使用x_train和一个下划线_来接受torch.sort的返回值，因为该函数会返回排序后的数据以及排序之前的下标。
• f(x)定义映射函数，即$y_i = 2\\sin(x_i) + x_i^0.8 + \\epsilon$
• y_train是使用f(x)生成训练数据的结果y。

生成测试集。

• 使用torch.arange生成$[0,5)$之间的数，步长为0.1
• f(x)生成测试集的真实结果
• n_test存储测试集数据的数量，结果也是50个。

def plot_kernel_reg(y_hat):
    d2l.plot(x_test, [y_truth, y_hat], x, y, legend=[Truth, Pred],
             xlim=[0, 5], ylim=[-1, 5])
    d2l.plt.plot(x_train, y_train, o, alpha=0.5);

用于画图的一个函数，不用深究什么意思。

# `X_repeat` 的形状: (`n_test`, `n_train`), 
# 每一行都包含着相同的测试输入（例如：同样的查询）
X_repeat = x_test.repeat_interleave(n_train).reshape((-1, n_train))
# `x_train` 包含着键。`attention_weights` 的形状：(`n_test`, `n_train`), 
# 每一行都包含着要在给定的每个查询的值（`y_train`）之间分配的注意力权重
attention_weights = nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)
# `y_hat` 的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = torch.matmul(attention_weights, y_train)
plot_kernel_reg(y_hat)

这里是使用非参的nadaraya-waston核回归计算attention权重，即$f(x) = \\sum_i=1^n \\fracK(x - xi)\\sumj=1^n K(x - x_j) y_i$这个公式可以根据x的位置对y进行加权。

• x_test.repeat_interleave(n_train)是将测试数据x_test重复n_train次即50次，再使用.reshape((-1, n_train))将其形状重置成50行50列。操作之后X_repeat是一个矩阵，其中一行50个数都是一样的，每列都是x_test

  >>tensor([[0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.1000, 0.1000, 0.1000,  ..., 0.1000, 0.1000, 0.1000],
          [0.2000, 0.2000, 0.2000,  ..., 0.2000, 0.2000, 0.2000],
          ...,
          [4.7000, 4.7000, 4.7000,  ..., 4.7000, 4.7000, 4.7000],
          [4.8000, 4.8000, 4.8000,  ..., 4.8000, 4.8000, 4.8000],
          [4.9000, 4.9000, 4.9000,  ..., 4.9000, 4.9000, 4.9000]])

• attention_weight那里就是将数据丢进一个softmax层里，要计算的核$K(x - x_i)$就是之前选定的高斯核$-\\frac12(x - x_i)^2$。使用广播机制进行计算，在这里：

  • x_train相当于key
  • X_repeat相当于query
  • y_train相当于value
• 将计算出来的attention_weight和y_train进行计算，即计算$\\sum_i=1^n \\mathrmsoftmax\\left(-\\frac12(x - x_i)^2\\right) y_i$
• 最终使用plot_kernel_reg将结果画出来。蓝色的线是真实数据，紫色虚线是我们模型预测的结果，可以看出：在数据量小的时候使用不带参数的核回归，虽然已经能看出大致曲线了但是效果还是差很多。
  

原书代码有两段可视化代码，这里只讲一段。另一端是heat map，表示注意力效果的。

class NWKernelRegression(nn.Module):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.w = nn.Parameter(torch.rand((1,), requires_grad=True))

    def forward(self, queries, keys, values):
        queries = queries.repeat_interleave(keys.shape[1]).reshape((-1, keys.shape[1]))
        self.attention_weights = nn.functional.softmax(-((queries - keys) * self.w)**2 / 2, dim=1)
        return torch.bmm(self.attention_weights.unsqueeze(1),values.unsqueeze(-1)).reshape(-1)

这里是带参的nadaraya-waston核回归，用的是$\\beginalignedf(x) &= \\sum_i=1^n \\mathrmsoftmax\\left(-\\frac12((x - x_i)w)^2\\right) y_i\\endaligned$了。

• forward定义前向传播，需要传入Q、K、V。

  • queries这里计算方法和非参的计算方法一样，将其重复复制成一个矩阵，为后边的softmax计算做准备。
  • self.attention_weights这里是$\\frac12((x - x_i)w)^2$，高斯核加上了参数，w的作用相当于控制高斯核的大小，可以想象成CNN中控制卷积核的大小。

  • 最后是用bmm计算乘法。这里涉及到一个mini-batch的矩阵乘法。

    补充bmm矩阵计算：使用bmm计算mini-batch矩阵乘法，需要三个参数，第一个参数是批量的数量，剩下两个参数是矩阵的维度。比如：

    X = torch.ones((2, 1, 4))
    Y = torch.ones((2, 4, 6))
    torch.bmm(X, Y).shape

    输出结果是

    torch.Size([2, 1, 6])

    • x是两个1*4的矩阵，y是两个4*的矩阵，使用bmm相乘之后打印一下结果的维度，是两个1*6的矩阵。

      weights = torch.ones((2, 10)) * 0.1
      values = torch.arange(20.0).reshape((2, 10))
      torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1))

      输出结果为：

      
      tensor([[[ 4.5000]],

    [[14.5000]]])

    
    - 使用小批量矩阵乘法来计算小批量数据中的加权平均值
    - unsqueeze是给张量添加维度的，具体可以看→[torch.squeeze 和 torch.unsqueeze](https://blog.51cto.com/Lolitann/5111261)
    - 这一步是假设w和v相乘，维度变化之后是2\\*1\\*10的矩阵和2\\*10\\*1的矩阵相乘，结果是2\\*1\\*1的矩阵，即两个batch。

X_tile = x_train.repeat((n_train, 1))
Y_tile = y_train.repeat((n_train, 1))
keys = X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))
values = Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))

将训练集数据转换成key和value。

• X_tile是将训练集的数据x_train平铺展开，每一行都包含着相同的训练输入。
• Y_tile是将训练集的数据y_train平铺展开，每一行都包含着相同的训练输出。

• torch.eye生成一个维度为n_train的单位矩阵。因为单位矩阵是对角线元素为1其余元素都为0 ，之后使用1 - torch.eye(n_train)将其转化为对角线元素为0其余元素都为1的矩阵，在进行类型转换，type(torch.bool)，使其转换为对角元素为false其余元素都为true的矩阵。

  key和value的这一步是将对角线上的数据mask掉。再将其转化为50行的矩阵，现在就变为50行49列的矩阵了（原来是50行50列）。
  
  画个单间的图你们理解一下：
  

net = NWKernelRegression()
loss = nn.MSELoss(reduction=none)
trainer = torch.optim.SGD(net.parameters(), lr=0.5)
animator = d2l.Animator(xlabel=epoch, ylabel=loss, xlim=[1, 5])

for epoch in range(5):
    trainer.zero_grad()
    # 注意：L2 Loss = 1/2 * MSE Loss。
    # PyTorch 的 MSE Loss 与 MXNet 的 L2Loss 差一个 2 的因子，因此被除2。
    l = loss(net(x_train, keys, values), y_train) / 2
    l.sum().backward()
    trainer.step()
    print(fepoch epoch + 1, loss float(l.sum()):.6f)
    animator.add(epoch + 1, float(l.sum()))

训练过程：

• 使用我们的带参数的计算方法
• loss使用MSE loss
• 优化器SGD
• d2l.Animator一个梯度下降过程的可视化，不用深究
• 之后就是训练5个epoch

keys = x_train.repeat((n_test, 1))
values = y_train.repeat((n_test, 1))
y_hat = net(x_test, keys, values).unsqueeze(1).detach()
plot_kernel_reg(y_hat)

训练之后再将结果进行可视化。拟合效果比不带参数的变好了，但是在注意力较大的区域参数变得不平滑。侧面也展示了注意力权重的影响。