论文:Attention Is All You Need
https://arxiv.org/pdf/1706.03762.pdf
解读:
https://blog.csdn.net/qq_41664845/article/details/84969266
https://blog.csdn.net/songbinxu/article/details/80332992
https://www.cnblogs.com/rosyYY/p/10115424.html
代码:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from .sublayers import MultiHeadAttention, PositionwiseFeedForward
from ..normalization import LayerNorm
class EncoderLayer(nn.Module):
"""
Represents one Encoder layer of the Transformer Encoder
Refer Fig. 1 in https://arxiv.org/pdf/1706.03762.pdf
NOTE: The layer normalization step has been moved to the input as per latest version of T2T
"""
def __init__(self, hidden_size, total_key_depth, total_value_depth, filter_size, num_heads,
bias_mask=None, layer_dropout=0.0, attention_dropout=0.0, relu_dropout=0.0):
"""
Parameters:
hidden_size: Hidden size
total_key_depth: Size of last dimension of keys. Must be divisible by num_head
total_value_depth: Size of last dimension of values. Must be divisible by num_head
output_depth: Size last dimension of the final output
filter_size: Hidden size of the middle layer in FFN
num_heads: Number of attention heads
bias_mask: Masking tensor to prevent connections to future elements
layer_dropout: Dropout for this layer
attention_dropout: Dropout probability after attention (Should be non-zero only during training)
relu_dropout: Dropout probability after relu in FFN (Should be non-zero only during training)
"""
super(EncoderLayer, self).__init__()
self.multi_head_attention = MultiHeadAttention(hidden_size, total_key_depth, total_value_depth,
hidden_size, num_heads, bias_mask, attention_dropout)
self.positionwise_feed_forward = PositionwiseFeedForward(hidden_size, filter_size, hidden_size,
layer_config='cc', padding = 'both',
dropout=relu_dropout)
self.dropout = nn.Dropout(layer_dropout)
self.layer_norm_mha = LayerNorm(hidden_size)
self.layer_norm_ffn = LayerNorm(hidden_size)
def forward(self, inputs):
x = inputs
# Layer Normalization
x_norm = self.layer_norm_mha(x)
# Multi-head attention
y = self.multi_head_attention(x_norm, x_norm, x_norm)
# Dropout and residual
x = self.dropout(x + y)
# Layer Normalization
x_norm = self.layer_norm_ffn(x)
# Positionwise Feedforward
y = self.positionwise_feed_forward(x_norm)
# Dropout and residual
y = self.dropout(x + y)
return y
class DecoderLayer(nn.Module):
"""
Represents one Decoder layer of the Transformer Decoder
Refer Fig. 1 in https://arxiv.org/pdf/1706.03762.pdf
NOTE: The layer normalization step has been moved to the input as per latest version of T2T
"""
def __init__(self, hidden_size, total_key_depth, total_value_depth, filter_size, num_heads,
bias_mask, layer_dropout=0.0, attention_dropout=0.0, relu_dropout=0.0):
"""
Parameters:
hidden_size: Hidden size
total_key_depth: Size of last dimension of keys. Must be divisible by num_head
total_value_depth: Size of last dimension of values. Must be divisible by num_head
output_depth: Size last dimension of the final output
filter_size: Hidden size of the middle layer in FFN
num_heads: Number of attention heads
bias_mask: Masking tensor to prevent connections to future elements
layer_dropout: Dropout for this layer
attention_dropout: Dropout probability after attention (Should be non-zero only during training)
relu_dropout: Dropout probability after relu in FFN (Should be non-zero only during training)
"""
super(DecoderLayer, self).__init__()
self.multi_head_attention_dec = MultiHeadAttention(hidden_size, total_key_depth, total_value_depth,
hidden_size, num_heads, bias_mask, attention_dropout)
self.multi_head_attention_enc_dec = MultiHeadAttention(hidden_size, total_key_depth, total_value_depth,
hidden_size, num_heads, None, attention_dropout)
self.positionwise_feed_forward = PositionwiseFeedForward(hidden_size, filter_size, hidden_size,
layer_config='cc', padding = 'left',
dropout=relu_dropout)
self.dropout = nn.Dropout(layer_dropout)
self.layer_norm_mha_dec = LayerNorm(hidden_size)
self.layer_norm_mha_enc = LayerNorm(hidden_size)
self.layer_norm_ffn = LayerNorm(hidden_size)
def forward(self, inputs):
"""
NOTE: Inputs is a tuple consisting of decoder inputs and encoder output
"""
x, encoder_outputs = inputs
# Layer Normalization before decoder self attention
x_norm = self.layer_norm_mha_dec(x)
# Masked Multi-head attention
y = self.multi_head_attention_dec(x_norm, x_norm, x_norm)
# Dropout and residual after self-attention
x = self.dropout(x + y)
# Layer Normalization before encoder-decoder attention
x_norm = self.layer_norm_mha_enc(x)
# Multi-head encoder-decoder attention
y = self.multi_head_attention_enc_dec(x_norm, encoder_outputs, encoder_outputs)
# Dropout and residual after encoder-decoder attention
x = self.dropout(x + y)
# Layer Normalization
x_norm = self.layer_norm_ffn(x)
# Positionwise Feedforward
y = self.positionwise_feed_forward(x_norm)
# Dropout and residual after positionwise feed forward layer
y = self.dropout(x + y)
# Return encoder outputs as well to work with nn.Sequential
return y, encoder_outputs
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class MultiHeadAttention(nn.Module):
"""
Multi-head attention as per https://arxiv.org/pdf/1706.03762.pdf
Refer Figure 2
"""
def __init__(self, input_depth, total_key_depth, total_value_depth, output_depth,
num_heads, bias_mask=None, dropout=0.0):
"""
Parameters:
input_depth: Size of last dimension of input
total_key_depth: Size of last dimension of keys. Must be divisible by num_head
total_value_depth: Size of last dimension of values. Must be divisible by num_head
output_depth: Size last dimension of the final output
num_heads: Number of attention heads
bias_mask: Masking tensor to prevent connections to future elements
dropout: Dropout probability (Should be non-zero only during training)
"""
super(MultiHeadAttention, self).__init__()
# Checks borrowed from
# https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/layers/common_attention.py
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_value_depth, num_heads))
self.num_heads = num_heads
self.query_scale = (total_key_depth//num_heads)**-0.5
self.bias_mask = bias_mask
# Key and query depth will be same
self.query_linear = nn.Linear(input_depth, total_key_depth, bias=False)
self.key_linear = nn.Linear(input_depth, total_key_depth, bias=False)
self.value_linear = nn.Linear(input_depth, total_value_depth, bias=False)
self.output_linear = nn.Linear(total_value_depth, output_depth, bias=False)
self.dropout = nn.Dropout(dropout)
def _split_heads(self, x):
"""
Split x such to add an extra num_heads dimension
Input:
x: a Tensor with shape [batch_size, seq_length, depth]
Returns:
A Tensor with shape [batch_size, num_heads, seq_length, depth/num_heads]
"""
if len(x.shape) != 3:
raise ValueError("x must have rank 3")
shape = x.shape
return x.view(shape[0], shape[1], self.num_heads, shape[2]//self.num_heads).permute(0, 2, 1, 3)
def _merge_heads(self, x):
"""
Merge the extra num_heads into the last dimension
Input:
x: a Tensor with shape [batch_size, num_heads, seq_length, depth/num_heads]
Returns:
A Tensor with shape [batch_size, seq_length, depth]
"""
if len(x.shape) != 4:
raise ValueError("x must have rank 4")
shape = x.shape
return x.permute(0, 2, 1, 3).contiguous().view(shape[0], shape[2], shape[3]*self.num_heads)
def forward(self, queries, keys, values):
# Do a linear for each component
queries = self.query_linear(queries)
keys = self.key_linear(keys)
values = self.value_linear(values)
# Split into multiple heads
queries = self._split_heads(queries)
keys = self._split_heads(keys)
values = self._split_heads(values)
# Scale queries
queries *= self.query_scale
# Combine queries and keys
logits = torch.matmul(queries, keys.permute(0, 1, 3, 2))
# Add bias to mask future values
if self.bias_mask is not None:
logits += self.bias_mask[:, :, :logits.shape[-2], :logits.shape[-1]].type_as(logits.data)
# Convert to probabilites
weights = nn.functional.softmax(logits, dim=-1)
# Dropout
weights = self.dropout(weights)
# Combine with values to get context
contexts = torch.matmul(weights, values)
# Merge heads
contexts = self._merge_heads(contexts)
#contexts = torch.tanh(contexts)
# Linear to get output
outputs = self.output_linear(contexts)
return outputs
class Conv(nn.Module):
"""
Convenience class that does padding and convolution for inputs in the format
[batch_size, sequence length, hidden size]
"""
def __init__(self, input_size, output_size, kernel_size, pad_type):
"""
Parameters:
input_size: Input feature size
output_size: Output feature size
kernel_size: Kernel width
pad_type: left -> pad on the left side (to mask future data),
both -> pad on both sides
"""
super(Conv, self).__init__()
padding = (kernel_size - 1, 0) if pad_type == 'left' else (kernel_size//2, (kernel_size - 1)//2)
self.pad = nn.ConstantPad1d(padding, 0)
self.conv = nn.Conv1d(input_size, output_size, kernel_size=kernel_size, padding=0)
def forward(self, inputs):
inputs = self.pad(inputs.permute(0, 2, 1))
outputs = self.conv(inputs).permute(0, 2, 1)
return outputs
class PositionwiseFeedForward(nn.Module):
"""
Does a Linear + RELU + Linear on each of the timesteps
"""
def __init__(self, input_depth, filter_size, output_depth, layer_config='ll', padding='left', dropout=0.0):
"""
Parameters:
input_depth: Size of last dimension of input
filter_size: Hidden size of the middle layer
output_depth: Size last dimension of the final output
layer_config: ll -> linear + ReLU + linear
cc -> conv + ReLU + conv etc.
padding: left -> pad on the left side (to mask future data),
both -> pad on both sides
dropout: Dropout probability (Should be non-zero only during training)
"""
super(PositionwiseFeedForward, self).__init__()
layers = []
sizes = ([(input_depth, filter_size)] +
[(filter_size, filter_size)]*(len(layer_config)-2) +
[(filter_size, output_depth)])
for lc, s in zip(list(layer_config), sizes):
if lc == 'l':
layers.append(nn.Linear(*s))
elif lc == 'c':
layers.append(Conv(*s, kernel_size=3, pad_type=padding))
else:
raise ValueError("Unknown layer type {}".format(lc))
self.layers = nn.ModuleList(layers)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(dropout)
def forward(self, inputs):
x = inputs
for i, layer in enumerate(self.layers):
x = layer(x)
if i < len(self.layers):
x = self.relu(x)
x = self.dropout(x)
return x