1. 首先加载包:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time
from torch.autograd import Variable
import matplotlib.pyplot as plt
import seaborn
seaborn.set_context(context="talk")
2. 常用的编码器 - 解码器结构、文本生成结构如下:
encode-函数,encoder-网络结构
src-输入文本,src_mask-掩码的输入文本,src_embed-嵌入后的输入文本
tgt-目标文本,tgt_mask-掩码的目标文本,tgt_embed-嵌入后的目标文本
memory-记忆
编码 encode:encoder(self.src_embed(src), src_mask)
解码 decode:decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
向前传播 forward:decode(self.encode(src, src_mask), src_mask,tgt, tgt_mask)
class EncoderDecoder(nn.Module):
"""
A standard Encoder-Decoder architecture. Base for this and many
other models.
"""
def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
super(EncoderDecoder, self).__init__()
self.encoder = encoder
self.decoder = decoder
self.src_embed = src_embed
self.tgt_embed = tgt_embed
self.generator = generator
def forward(self, src, tgt, src_mask, tgt_mask):
"Take in and process masked src and target sequences."
return self.decode(self.encode(src, src_mask), src_mask,
tgt, tgt_mask)
def encode(self, src, src_mask):
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
文本生成结构如下:
线性层:nn.Linear(d_model, vocab)
Softmax层:F.log_softmax(self.proj(x), dim=-1)
class Generator(nn.Module):
"Define standard linear + softmax generation step."
def __init__(self, d_model, vocab):
super(Generator, self).__init__()
self.proj = nn.Linear(d_model, vocab)
def forward(self, x):
return F.log_softmax(self.proj(x), dim=-1)
3. 编码器
首先定义编码器,六个相同的层采用深拷贝搭建:
def clones(module, N):
# 产生N个相同的层,N=6
# ModuleList 可以像常规Python列表一样编制索引,包含的模块已正确注册
# copy.copy 浅拷贝 只拷贝父对象,不会拷贝对象的内部的子对象
# copy.deepcopy 深拷贝 拷贝对象及其子对象
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
class Encoder(nn.Module):
# "Core encoder is a stack of N layers"
def __init__(self, layer, N):
super(Encoder, self).__init__()
self.layers = clones(layer, N)
# 归一化层 LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True)
# normalized_shape 输入尺寸 [∗×normalized_shape[0]×normalized_shape[1]×…×normalized_shape[−1]]
# eps-为保证数值稳定性(分母不能趋近或取0),给分母加上的值。默认为1e-5
# elementwise_affine 布尔值,当设为true,给该层添加可学习的仿射变换参数
self.norm = LayerNorm(layer.size)
def forward(self, x, mask):
# "Pass the input (and mask) through each layer in turn."
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
然后构建LayerNorm,在两个子层中分别使用残余连接,然后是层标准化
class LayerNorm(nn.Module):
def __init__(self, features, eps=1e-6):
super(LayerNorm, self).__init__()
self.a_2 = nn.Parameter(torch.ones(features))
self.b_2 = nn.Parameter(torch.zeros(features))
self.eps = eps
def forward(self, x):
mean = x.mean(-1, keepdim=True)
std = x.std(-1, keepdim=True)
return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
每个子层的输出是,其中是由子层本身实现的函数。将dropout应用于每个子层的输出,然后将其添加到子层输入并进行规范化。为了促进这些残余连接,模型中的所有子层以及嵌入层都产生维度为 的输出。
class SublayerConnection(nn.Module):
"""
A residual connection followed by a layer norm.
Note for code simplicity the norm is first as opposed to last.
"""
def __init__(self, size, dropout):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(size)
self.dropout = nn.Dropout(dropout)
def forward(self, x, sublayer):
"Apply residual connection to any sublayer with the same size."
return x + self.dropout(sublayer(self.norm(x)))
每层有两个子层组成,一是多头自我关注机制multi-head self-attention mechanism,即实现了“多头”的 Self-attention,二是位置完全连接的前馈网络position-wise fully connected feed- forward network。
class EncoderLayer(nn.Module):
"Encoder is made up of self-attn and feed forward (defined below)"
def __init__(self, size, self_attn, feed_forward, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = self_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 2)
self.size = size
def forward(self, x, mask):
"Follow Figure 1 (left) for connections."
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
return self.sublayer[1](x, self.feed_forward)
4. 解码器
新增加了一个编码-解码多头注意力层,其他与编码器相同,也是由6个相同层组成:
class Decoder(nn.Module):
"Generic N layer decoder with masking."
# N=6
def __init__(self, layer, N):
super(Decoder, self).__init__()
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def forward(self, x, memory, src_mask, tgt_mask):
for layer in self.layers:
x = la