import math import torch from torch import nn from d2l import torch as d2l class PositionalEncoding(nn.Module): def __init__(self, num_hiddens, dropout, max_len=1000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(dropout) self.P = torch.zeros((1, max_len, num_hiddens)) X = torch.arange(max_len, dtype=torch.float32).reshape( -1, 1) / torch.pow(10000, torch.arange( 0, num_hiddens, 2, dtype=torch.float32) / num_hiddens) self.P[:, :, 0::2] = torch.sin(X) self.P[:, :, 1::2] = torch.cos(X) def forward(self, X): X = X + self.P[:, :X.shape[1], :].to(X.device) return self.dropout(X) encoding_dim, num_steps = 32, 60 pos_encoding = PositionalEncoding(encoding_dim, 0) pos_encoding.eval() X = pos_encoding(torch.zeros((1, num_steps, encoding_dim))) P = pos_encoding.P[:, :X.shape[1], :] d2l.plot(torch.arange(num_steps), P[0, :, 6:10].T, xlabel=’Row (position)’, figsize=(6, 2.5), legend=["Col %d" % d for d in torch.arange(6, 10)])
时间: 2024-04-19 13:24:13 浏览: 10
这段代码演示了位置编码(Positional Encoding)的计算过程。
首先,定义了一个名为 `PositionalEncoding` 的类,继承自 `nn.Module`。在初始化方法 `__init__` 中,定义了位置编码所需的参数,包括隐藏层大小 `num_hiddens`、dropout 比例 `dropout` 和最大序列长度 `max_len`。然后,创建了一个形状为 `(1, max_len, num_hiddens)` 的张量 `self.P`,用于存储位置编码的值。
接下来,通过计算正弦和余弦函数的值,将位置信息编码到张量 `self.P` 中。具体地,首先创建一个形状为 `(max_len, 1)` 的张量 `X`,其中元素的值从 0 到 `max_len-1`。然后,根据隐藏层大小 `num_hiddens`,计算每个位置上的编码值,并将其分别赋值给 `self.P` 张量的奇数列和偶数列。
在 `forward` 方法中,输入张量 `X` 与位置编码张量 `self.P` 相加,并返回结果。同时,还通过 `dropout` 层对结果进行随机失活处理。
接着,创建了一个位置编码实例 `pos_encoding`,传入隐藏层大小和 dropout 比例,并将其设为评估模式。然后,创建一个形状为 `(1, num_steps, encoding_dim)` 的全零张量 `X`,并将其传入位置编码实例中进行编码。最后,绘制了位置编码张量的部分列的曲线图。
需要注意的是,在绘制曲线图之前,对位置编码实例调用了 `eval()` 方法,将其设为评估模式,以避免在推理过程中应用随机失活。