| 6 | from torch import Tensor |
| 7 | |
| 8 | class PositionalEncoding(nn.Module): |
| 9 | def __init__(self, max_positions, dim_embed, drop_prob): |
| 10 | super().__init__() |
| 11 | |
| 12 | assert dim_embed % 2 == 0 |
| 13 | |
| 14 | # 生成一个max_positions行一列, 所有值在0到max_positions-1之间的矩阵 |
| 15 | position = torch.arange(max_positions).unsqueeze(1) |
| 16 | dim_pair = torch.arange(0, dim_embed, 2) |
| 17 | div_term = torch.exp(dim_pair * (-math.log(10000.0) / dim_embed)) |
| 18 | |
| 19 | pe = torch.zeros(max_positions, dim_embed) |
| 20 | pe[:, 0::2] = torch.sin(position * div_term) # 取出偶数位元素 |
| 21 | # print('pe[:, 0::2]:', pe[:, 0::2][0]) |
| 22 | # print('pe[:, 0::2].shape:', pe[:, 0::2].shape) # torch.Size([100, 256]) |
| 23 | pe[:, 1::2] = torch.cos(position * div_term) # 取出奇数位元素 |
| 24 | |
| 25 | # 扩充batch维度 |
| 26 | pe = pe.unsqueeze(0) |
| 27 | |
| 28 | # 整个学习阶段，位置信息是不变的，固定为不可学习的数据 |
| 29 | self.register_buffer('pe', pe) |
| 30 | self.dropout = nn.Dropout(p=drop_prob) |
| 31 | |
| 32 | def forward(self, x): |
| 33 | # 计算每个batch的最大句子长度 |
| 34 | max_squence_length = x.size(1) |
| 35 | |
| 36 | # 词向量中添加位置信息 |
| 37 | x = x + self.pe[:, :max_squence_length] |
| 38 | x = self.dropout(x) |
| 39 | return x |
| 40 | |
| 41 | |
| 42 | if __name__ == "__main__": |