La última traducción de The Annotated Transformer para la edición 2023
Dirección original: http://nlp.seas.harvard.edu/annotated-transformer/#hardware-and-schedule
0.Preliminares
import os
from os.path import exists
import torch
import torch.nn as nn
from torch.nn.functional import log_softmax, pad
import math
import copy
import time
from torch.optim.lr_scheduler import LambdaLR
import pandas as pd
import altair as alt
from torchtext.data.functional import to_map_style_dataset
from torch.utils.data import DataLoader
from torchtext.vocab import build_vocab_from_iterator
import torchtext.datasets as datasets
import spacy
import GPUtil
import warnings
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
# Set to False to skip notebook execution (e.g. for debugging)
warnings.filterwarnings("ignore")
RUN_EXAMPLES = True
algunas funciones de ayuda pública
def is_interactive_notebook():
return __name__ == "__main__"
def show_example(fn, args=[]):
if __name__ == "__main__" and RUN_EXAMPLES:
return fn(*args)
def execute_example(fn, args=[]):
if __name__ == "__main__" and RUN_EXAMPLES:
fn(*args)
class DummyOptimizer(torch.optim.Optimizer):
def __init__(self):
self.param_groups = [{"lr": 0}]
None
def step(self):
None
def zero_grad(self, set_to_none=False):
None
class DummyScheduler:
def step(self):
None
1. Antecedentes
La GPU neuronal extendida, ByteNet y ConvS2S surgieron para reducir la cantidad de cómputo secuencial, y todos usan redes neuronales convolucionales como bloques de construcción básicos para calcular representaciones ocultas de todas las posiciones de entrada y salida en paralelo. En estos modelos, la cantidad de operaciones requeridas para correlacionar señales de dos ubicaciones de entrada o salida arbitrarias crece linealmente para ConvS2S y logarítmicamente para ByteNet con la distancia entre las ubicaciones. Esto hace que sea más difícil aprender las dependencias entre ubicaciones distantes. En el Transformador, esto se reduce a un número constante de operaciones, aunque con un efecto reducido debido a la posición ponderada de atención promedio, que contrarrestamos con atención de múltiples cabezas.
La autoatención, a veces llamada intraatención, es un mecanismo de atención que asocia diferentes posiciones de una sola secuencia para calcular representaciones de secuencias. La autoatención se ha utilizado con éxito para varias tareas, incluida la comprensión de lectura, el resumen abstracto, la vinculación textual y el aprendizaje de representaciones de oraciones independientes de la tarea. Las redes de memoria de extremo a extremo se basan en mecanismos de atención recurrentes en lugar de ciclos de alineación de secuencias, y se ha demostrado que funcionan bien en tareas de modelado de lenguaje y respuesta a preguntas de lenguaje simple.
Sin embargo, hasta donde sabemos, Transformer es el primer modelo que se basa completamente en la autoatención para calcular sus representaciones de entrada y salida, en lugar de utilizar convoluciones o RNN alineados en secuencia.
2. Arquitectura modelo
La mayoría de los modelos de traducción de secuencias neuronales tienen una estructura de codificador-decodificador . El codificador toma una secuencia de entrada ( x 1 , . . . , xn ) (x_1, ..., x_n)( X1,... ,Xn) se asigna a una representación continua $z=(z_1, ..., z_n). El decodificador produce una secuencia de salida para cada elemento en z. El decodificador produce una secuencia de salida para cada elemento en zmedio. El decodificador genera una secuencia de salida (y_1, ..., y_m)$ para cada elemento en z y genera un elemento a la vez. En cada paso, el modelo esautorregresivo, agregando estructuras generadas previamente a la secuencia de entrada para predecir juntos cuándo generar el siguiente resultado. (Características de los modelos autorregresivos)
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)
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 log_softmax(self.proj(x), dim=-1)
Tanto el codificador como el decodificador de Transformer utilizan apilamiento de atención automática y capas completamente conectadas por puntos. Como se muestra en los lados izquierdo y derecho de la Figura 1.
2.1 Codificador y Decodificador
El codificador consta de N = 6 capas idénticas.
(1) Codificador
def clones(module, N):
"Produce N identical layers."
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)
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)
Cada subcapa del codificador (capa de autoatención y FFNN) está conectada a una conexión residual (cite) . Entonces capa-normalización (cite) .
class LayerNorm(nn.Module):
"Construct a layernorm module (See citation for details)."
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
La salida de cada subcapa es LayerNorm ( x + Sublayer ( x ) ) LayerNorm(x + Sublayer(x))Capa N o m ( x _ _+S u b l a yer ( x )) , donde $Sublayer(x) $ es una función implementada por la propia subcapa. Aplicamos dropout(cite)a la salida de cada subcapa antes de agregarla a la entrada de la subcapa y normalizarla.
Para facilitar la conexión residual, las dimensiones de todas las subcapas en el modelo y la salida generada por la capa incrustada son dmodel = 512 d_{model}=512dm o d e l=512
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)))
Cada capa tiene dos subcapas. La primera capa es un mecanismo de autoatención de varios cabezales (capa), y la segunda capa es una red de retroalimentación simple y totalmente conectada.
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)
(2) Decodificador
El decodificador también consta de N = 6 capas idénticas.
class Decoder(nn.Module):
"Generic N layer decoder with masking."
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 = layer(x, memory, src_mask, tgt_mask)
return self.norm(x)
Además de las dos subcapas en cada capa del decodificador, el decodificador tiene una tercera subcapa que realiza atención multicabezal en la salida del codificador. (Es decir, la capa codificador-decodificador-atención, el vector q proviene de la entrada de la capa anterior, y los vectores k y v son la memoria vectorial de salida de la última capa del codificador) Similar al codificador, usamos conexiones residuales en cada subcapa, y luego realice la normalización de la capa.
class DecoderLayer(nn.Module):
"Decoder is made of self-attn, src-attn, and feed forward (defined below)"
def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
super(DecoderLayer, self).__init__()
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 3)
def forward(self, x, memory, src_mask, tgt_mask):
"Follow Figure 1 (right) for connections."
m = memory
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
return self.sublayer[2](x, self.feed_forward)
También modificamos la subcapa de autoatención en la capa del decodificador para evitar que la posición actual preste atención a la posición posterior. Esta combinación de máscaras compensa la incrustación de salida en una posición, asegurando que la posición iiLa predicción de i depende solo de la posicióniiSalida conocida antes de i .
def subsequent_mask(size):
"Mask out subsequent positions."
attn_shape = (1, size, size)
subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)
return subsequent_mask == 0
La máscara de atención a continuación muestra dónde se le permite buscar cada palabra tgt (fila) (columna). Durante el entrenamiento, la información futura de la palabra actual se bloquea para evitar que esta palabra preste atención a las siguientes palabras.
def example_mask():
LS_data = pd.concat(
[
pd.DataFrame(
{
"Subsequent Mask": subsequent_mask(20)[0][x, y].flatten(),
"Window": y,
"Masking": x,
}
)
for y in range(20)
for x in range(20)
]
)
return (
alt.Chart(LS_data)
.mark_rect()
.properties(height=250, width=250)
.encode(
alt.X("Window:O"),
alt.Y("Masking:O"),
alt.Color("Subsequent Mask:Q", scale=alt.Scale(scheme="viridis")),
)
.interactive()
)
show_example(example_mask)
(3) Atención
La función Atención se puede describir como la asignación de una consulta y un conjunto de pares clave-valor a una salida, donde la consulta, la clave, el valor y la salida son todos vectores. La salida es la suma ponderada de valores, donde el peso de cada valor se calcula mediante la función compatible de la consulta y la clave correspondiente.
A la atención particular la llamamos "Scaled Dot-Product Attention" (Atención de producto escalado escalado"). Su entrada es query, key (dimensión es dk d_kdk) y valores (dimensión es dv d_vdv). Calculamos el producto punto de consulta y todas las claves, y luego dividimos cada uno por dk \sqrt{d_k}dk, y finalmente use la función softmax para obtener el peso del valor.
En la práctica, calculamos simultáneamente funciones de atención para un conjunto de consultas y las combinamos en una matriz $Q . La clave y el valor también forman una matriz juntos. La clave y el valor también forman una matriz juntos. clave y valor también forman la matriz K y _ _y V$. La matriz de salida que calculamos es:
Atención ( Q , K , V ) = softmax ( QKT dk ) V Atención(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})VA t e n c i ó n ( Q ,K ,V )=entonces f t máx ( _dkQ KT) V
def attention(query, key, value, mask=None, dropout=None):
"Compute 'Scaled Dot Product Attention'"
d_k = query.size(-1)
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
p_attn = scores.softmax(dim=-1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
Las dos funciones de atención más utilizadas son la atención aditiva (citar) y la atención de producto escalar (multiplicación). Además del factor de escala 1 dk \frac{1}{\sqrt{d_k}}dk1, el producto escalar Atención es el mismo que nuestro algoritmo habitual. La atención aditiva calcula funciones compatibles utilizando una red de avance con una sola capa oculta. Aunque teóricamente la atención del producto punto y la atención aditiva tienen una complejidad similar, en la práctica, la atención del producto punto se puede implementar utilizando una multiplicación de matrices altamente optimizada, por lo que el cálculo de la atención del producto punto es más rápido y más eficiente en cuanto al espacio.
cuando dk d_kdkCuando el valor de es relativamente pequeño, el rendimiento de los dos mecanismos es similar. cuando dk d_kdkCuando es más grande, la atención aditiva funciona mejor que la atención de producto escalar sin escalar (citar). Sospechamos que para dk d_k muy grandedkvalor, el producto escalar crece sustancialmente, empujando la función softmax a regiones con gradientes extremadamente pequeños. (Para ilustrar por qué el producto escalar se vuelve más grande, suponga que qqq ykk es una variable aleatoria independiente con media 0 y varianza 1. Entonces su producto escalarq ⋅ k = ∑ i = 1 dkqikiq \cdot k = \sum_{i=1}^{d_k} q_ik_iq⋅k=∑yo = 1dkqyokyo, la media es 0 y la varianza es dk d_kdk). Para contrarrestar este efecto, reducimos el producto escalar en 1 dk \frac{1}{\sqrt{d_k}}dk1veces.
¿Por qué dividir por ** d \sqrt{d} en Atención ?d**** ¿muy importante? **
El cálculo de Atención es realizar softmax después del producto interno, y la operación principal involucrada es eq ⋅ ke^{q \cdot k}miq ⋅ k , podemos pensar aproximadamente que el valor después del producto interno y antes del softmax es− 3 d -3\sqrt{d}− 3da 3 d 3\sqrt{d}3dEn este rango, dado que d suele ser al menos 64, entonces e 3 de^{3\sqrt{d}}mi3dRelativamente grande y e − 3 de^{-3\sqrt{d}}mi− 3dEs relativamente pequeño, por lo que después de softmax, la distribución de Atención es muy cercana a una distribución caliente, lo que genera un problema grave de desaparición de gradiente, lo que resulta en un efecto de entrenamiento deficiente. (Por ejemplo, y=softmax(x) ingresa a la zona de saturación cuando |x| es grande, y el valor de y permanece casi sin cambios cuando x continúa cambiando, es decir, el gradiente en la zona de saturación desaparece). hay dos soluciones: 1.
Al igual que
NTK paramétricamente, dividir por d \sqrt{d} después del producto internod, de modo que la varianza de q⋅k se convierte en 1, correspondiente a e 3 e^3mi3 ,mi − 3 e^{−3}mi− 3 no es demasiado grande ni demasiado pequeño, por lo que después de softmax, no se convertirá en uno caliente y el gradiente desaparecerá. Esta es también la práctica de los transformadores convencionales como Self Attention en BERT 2. Además, no se divide
pord \sqrt{d }d, pero al inicializar la capa completamente conectada de q y k, la varianza de inicialización debe dividirse por una d más, lo que también puede hacer que la varianza inicial de q⋅k se convierta en 1. T5 adopta este enfoque.
La atención de múltiples cabezas permite que el modelo se enfoque conjuntamente en la información de diferentes subespacios de representación en diferentes ubicaciones, y si solo hay una cabeza de atención, su promedio debilitará esta información.
M ulti H ead ( Q , K , V ) = C oncat ( head 1 , . . . , headh ) WO donde headi = Atención ( QW i Q , KW i K , VW i V ) MultiHead(Q,K,V )=Concat(cabeza_1,...,cabeza_h)W^O \\ donde ~ cabeza_i = Atención(QW_i^Q, KW_i^K, VW_i^V)M u lt i Head a d ( Q ,K ,V )=C o n c a t ( cabeza _ _ _1,... ,cabeza _ _ _h) WOa dónde se dirige _ _ _ yo=A t e n c i ó n ( Q Wiq,KW _ik,VW _iV)
El mapeo se realiza mediante la matriz de peso: $W^Q_i \in \mathbb{R}^{d_{ {
modelo}} \times d_k}
$, W i K ∈ R d modelo × dk W^K_i \in \mathbb {R }^{d_{\text{modelo}} \times d_k}Wik∈Rdmodelo× rek, W i V ∈ R d modelo × dv W^V_i \in \mathbb{R}^{d_{\text{modelo}} \times d_v}WiV∈Rdmodelo× rev和W i O ∈ R hdv × d modelo W^O_i \in \mathbb{R}^{hd_v \times d_{\text{modelo}} }WiO∈Rhd _v× remodelo。
En este trabajo empleamos h = 8 capas o cabezas de atención paralelas. Para cada uno de estos cabezales, usamos $d_k=d_v=d_{\text{model}}/h=64 $, y debido a la reducida dimensionalidad de cada cabezal, el costo computacional total es similar a la atención de un solo cabezal con todos dimensiones
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
"Take in model size and number of heads."
super(MultiHeadedAttention, self).__init__()
assert d_model % h == 0
# We assume d_v always equals d_k
self.d_k = d_model // h
self.h = h
self.linears = clones(nn.Linear(d_model, d_model), 4)
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
"Implements Figure 2"
if mask is not None:
# Same mask applied to all h heads.
mask = mask.unsqueeze(1)
nbatches = query.size(0)
# 1) Do all the linear projections in batch from d_model => h x d_k
query, key, value = [
lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for lin, x in zip(self.linears, (query, key, value))
]
# 2) Apply attention on all the projected vectors in batch.
x, self.attn = attention(
query, key, value, mask=mask, dropout=self.dropout
)
# 3) "Concat" using a view and apply a final linear.
x = (
x.transpose(1, 2)
.contiguous()
.view(nbatches, -1, self.h * self.d_k)
)
del query
del key
del value
return self.linears[-1](x)
(4) Aplicación de la Atención en el modelo
Hay tres formas diferentes de usar la atención de múltiples cabezas en Transformer:
- En la capa de atención del codificador-decodificador, las consultas provienen de la capa del decodificador anterior y las claves y valores provienen de la salida del codificador. Esto permite que cada posición en el decodificador preste atención a todas las posiciones en la secuencia de entrada. Este es un mecanismo de atención que imita un codificador-descodificador típico en un modelo de secuencia a secuencia, como ( cite ).
- El codificador contiene una capa de autoatención. En la capa de autoatención, todas las claves, valores y consultas provienen del mismo lugar, que es la salida de la capa anterior en el codificador. En este caso, cada posición en el codificador puede prestar atención a todas las posiciones en la capa por encima del codificador.
- De manera similar, una capa de autoatención en el decodificador permite que cada posición en el decodificador preste atención a todas las posiciones que preceden a la posición actual en la capa del decodificador (incluida la posición actual). Para preservar las propiedades autorregresivas del decodificador, es necesario evitar que la información del decodificador fluya hacia la izquierda. Dentro de la atención del producto punto escalado, enmascaramos todos los valores de conexión ilegales en la entrada softmax (establecida en − ∞ -\infty− ∞ ) logra esto.
2.2 Red de retroalimentación basada en la posición
A excepción de la subcapa de atención, cada capa de nuestro codificador y decodificador contiene una red de retroalimentación completamente conectada que se aplica a cada ubicación por separado y de manera idéntica. La red consta de dos transformaciones lineales con una función de activación de ReLU en el medio.
FFN ( x ) = máx ( 0 , x W 1 + b 1 ) W 2 + b 2 FFN(x)=max(0, xW_1+b_1)W_2+b_2FFN ( x )=máx ( 0 , _x ancho1+b1) W2+b2
La posición es cada token en la secuencia,
Position-wisees decir, el MLP se aplica a cada token una vez y se aplica el mismo MLP.
Aunque ambas capas son transformaciones lineales, utilizan diferentes parámetros de una capa a otra. Otra forma de describirlo es como una convolución de dos núcleos de tamaño 1. Tanto las dimensiones de entrada como las de salida son $d_{\text{modelo}}=512 la dimensión interna es la dimensión interna esLa dimensión de la capa interna es d_{ff}=2048$ (es decir, la primera capa ingresa 512 dimensiones y genera 2048 dimensiones; la segunda capa ingresa 2048 dimensiones y genera 512 dimensiones)
class PositionwiseFeedForward(nn.Module):
"Implements FFN equation."
def __init__(self, d_model, d_ff, dropout=0.1):
super(PositionwiseFeedForward, self).__init__()
self.w_1 = nn.Linear(d_model, d_ff)
self.w_2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
return self.w_2(self.dropout(self.w_1(x).relu()))
2.3 Embebidos y Softmax
De manera similar a otros modelos de transformación de secuencias, usamos incrustaciones aprendidas para transformar tokens de entrada y tokens de salida en el modelo d d_{\text{modelo}}dmodelovector de dimensión También convertimos la salida del decodificador a la probabilidad predicha del próximo token usando una transformación lineal ordinaria y una función softmax. En nuestro modelo, la misma matriz de peso se comparte entre las dos capas de incrustación y la transformación lineal anterior a Softmax, similar a (citar). En la capa incrustada, multiplicamos estos pesos por $\sqrt{d_{\text{model}}} $
class Embeddings(nn.Module):
def __init__(self, d_model, vocab):
super(Embeddings, self).__init__()
self.lut = nn.Embedding(vocab, d_model)
self.d_model = d_model
def forward(self, x):
return self.lut(x) * math.sqrt(self.d_model)
2.4 Codificación de posición
Dado que nuestro modelo no contiene bucles ni circunvoluciones, para que el modelo aproveche el orden de la secuencia, debemos agregar información sobre la posición relativa o absoluta de los tokens en la secuencia . Para hacer esto, agregamos una "codificación posicional" a la incrustación de entrada en la parte inferior de la pila del codificador y del decodificador. Las dimensiones de la codificación de posición y la incrustación son las mismas, y también es d modelo d_{\text{modelo}}dmodelo, por lo que los dos vectores se pueden sumar. Hay varios códigos de posición para elegir, como códigos de posición aprendidos y códigos de posición fijos ( citar ).
En este trabajo utilizamos las funciones seno y coseno a diferentes frecuencias:
PE ( pos , 2 i ) = sin ( pos / 1000 0 2 i / dmodelo ) PE_{(pos,2i)}=sin(pos/10000^{2i/d_{modelo}})educación física( pos , 2 i ) _=sin(pos/100002 i / dm o d e l)
PE ( pos , 2 i + 1 ) = cos ( pos / 1000 0 2 i / dmodelo ) PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{modelo}})educación física( pos , 2 i + 1 ) _=cos ( pos / 1000 02 i / dm o d e l)
donde pos posp os es posición,iii es la dimensión. Es decir, cada dimensión del código de posición corresponde a una sinusoide. Estas longitudes de onda forman a de2π 2\pi2 π到10000 ⋅ 2 π 10000 \cdot 2\pi10000⋅La serie de conjuntos de 2 π . Elegimos esta función porque planteamos la hipótesis de que facilitaría que el modelo aprendiera a prestar atención a las posiciones relativas, ya que para cualquier desplazamiento k dado,PE pos + k PE_{pos+k}educación físicapos + k _Se puede expresar como PE pos PE_{pos}educación físicapos _función lineal de .
Además, agregamos un abandono a la suma de las codificaciones posicionales e incrustadas en las pilas de codificador y decodificador. Para el modelo base, la tasa de abandono que usamos es P drop = 0.1 P_{drop}=0.1PAGdejar caer _=0.1
class PositionalEncoding(nn.Module):
"Implement the PE function."
def __init__(self, d_model, dropout, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
# Compute the positional encodings once in log space.
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) *
-(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer("pe", pe)
def forward(self, x):
x = x + self.pe[:, : x.size(1)].requires_grad_(False)
return self.dropout(x)
La codificación de posición agregará una onda sinusoidal basada en la posición. La frecuencia y el desplazamiento de las ondas son diferentes para cada dimensión.
def example_positional():
pe = PositionalEncoding(20, 0)
y = pe.forward(torch.zeros(1, 100, 20))
data = pd.concat(
[
pd.DataFrame(
{
"embedding": y[0, :, dim],
"dimension": dim,
"position": list(range(100)),
}
)
for dim in [4, 5, 6, 7]
]
)
return (
alt.Chart(data)
.mark_line()
.properties(width=800)
.encode(x="position", y="embedding", color="dimension:N")
.interactive()
)
show_example(example_positional)
También intentamos usar incrustaciones posicionales aprendidas ( cite ) en su lugar, y descubrimos que ambas versiones produjeron resultados casi idénticos. Elegimos la versión sinusoidal porque permite que el modelo infiera longitudes de secuencia más largas que las encontradas durante el entrenamiento.
2.5 Modelo completo
Aquí definimos una función de hiperparámetros al modelo completo.
def make_model(
src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
):
"Helper: Construct a model from hyperparameters."
c = copy.deepcopy
attn = MultiHeadedAttention(h, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
model = EncoderDecoder(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
Generator(d_model, tgt_vocab),
)
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
return model
2.6 Inferencia
Aquí, usamos las predicciones del modelo generativo. Intentamos usar nuestro transformador para recordar la entrada. Como verá, dado que el modelo no ha sido entrenado, la salida se genera aleatoriamente. En el siguiente tutorial, construiremos la función de entrenamiento e intentaremos entrenar nuestro modelo para recordar los números del 1 al 10.
def inference_test():
test_model = make_model(11, 11, 2)
test_model.eval()
src = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
src_mask = torch.ones(1, 1, 10)
memory = test_model.encode(src, src_mask)
ys = torch.zeros(1, 1).type_as(src)
for i in range(9):
out = test_model.decode(
memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
)
prob = test_model.generator(out[:, -1])
_, next_word = torch.max(prob, dim=1)
next_word = next_word.data[0]
ys = torch.cat(
[ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1
)
print("Example Untrained Model Prediction:", ys)
def run_tests():
for _ in range(10):
inference_test()
show_example(run_tests)
Example Untrained Model Prediction: tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
Example Untrained Model Prediction: tensor([[0, 3, 4, 4, 4, 4, 4, 4, 4, 4]])
Example Untrained Model Prediction: tensor([[ 0, 10, 10, 10, 3, 2, 5, 7, 9, 6]])
Example Untrained Model Prediction: tensor([[ 0, 4, 3, 6, 10, 10, 2, 6, 2, 2]])
Example Untrained Model Prediction: tensor([[ 0, 9, 0, 1, 5, 10, 1, 5, 10, 6]])
Example Untrained Model Prediction: tensor([[ 0, 1, 5, 1, 10, 1, 10, 10, 10, 10]])
Example Untrained Model Prediction: tensor([[ 0, 1, 10, 9, 9, 9, 9, 9, 1, 5]])
Example Untrained Model Prediction: tensor([[ 0, 3, 1, 5, 10, 10, 10, 10, 10, 10]])
Example Untrained Model Prediction: tensor([[ 0, 3, 5, 10, 5, 10, 4, 2, 4, 2]])
Example Untrained Model Prediction: tensor([[0, 5, 6, 2, 5, 6, 2, 6, 2, 2]])
3. Entrenamiento
Esta sección describe el mecanismo de entrenamiento de nuestro modelo.
Aquí presentamos rápidamente algunas herramientas para entrenar un modelo estándar de codificador-decodificador. Primero, definimos un objeto por lotes que contiene las oraciones src y target para el entrenamiento, y construye la máscara.
3.1 Lotes y enmascaramiento
class Batch:
"""Object for holding a batch of data with mask during training."""
def __init__(self, src, tgt=None, pad=2): # 2 = <blank>
self.src = src
self.src_mask = (src != pad).unsqueeze(-2)
if tgt is not None:
self.tgt = tgt[:, :-1]
self.tgt_y = tgt[:, 1:]
self.tgt_mask = self.make_std_mask(self.tgt, pad)
self.ntokens = (self.tgt_y != pad).data.sum()
@staticmethod
def make_std_mask(tgt, pad):
"Create a mask to hide padding and future words."
tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(
tgt_mask.data
)
return tgt_mask
A continuación, creamos una función común de capacitación y evaluación para realizar un seguimiento de la pérdida. Pasamos una función de pérdida genérica, que también se usa para actualizaciones de parámetros.
3.2 Bucle de entrenamiento
class TrainState:
"""Track number of steps, examples, and tokens processed"""
step: int = 0 # Steps in the current epoch
accum_step: int = 0 # Number of gradient accumulation steps
samples: int = 0 # total # of examples used
tokens: int = 0 # total # of tokens processed
def run_epoch(
data_iter,
model,
loss_compute,
optimizer,
scheduler,
mode="train",
accum_iter=1,
train_state=TrainState(),
):
"""Train a single epoch"""
start = time.time()
total_tokens = 0
total_loss = 0
tokens = 0
n_accum = 0
for i, batch in enumerate(data_iter):
out = model.forward(
batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
)
loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)
# loss_node = loss_node / accum_iter
if mode == "train" or mode == "train+log":
loss_node.backward()
train_state.step += 1
train_state.samples += batch.src.shape[0]
train_state.tokens += batch.ntokens
if i % accum_iter == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
n_accum += 1
train_state.accum_step += 1
scheduler.step()
total_loss += loss
total_tokens += batch.ntokens
tokens += batch.ntokens
if i % 40 == 1 and (mode == "train" or mode == "train+log"):
lr = optimizer.param_groups[0]["lr"]
elapsed = time.time() - start
print(
(
"Epoch Step: %6d | Accumulation Step: %3d | Loss: %6.2f "
+ "| Tokens / Sec: %7.1f | Learning Rate: %6.1e"
)
% (i, n_accum, loss / batch.ntokens, tokens / elapsed, lr)
)
start = time.time()
tokens = 0
del loss
del loss_node
return total_loss / total_tokens, train_state
3.3 Datos de entrenamiento y procesamiento por lotes
Entrenamos en el conjunto de datos estándar WMT 2014 inglés-alemán que contiene alrededor de 4,5 millones de pares de oraciones. Las oraciones se codifican mediante pares de bytes, y las oraciones de origen y de destino comparten un vocabulario de aproximadamente 37 000 tokens. Para la traducción inglés-francés, utilizamos el conjunto de datos inglés-francés WMT 2014, significativamente más grande, que consta de 36 millones de oraciones y tokens divididos en 32 000 vocabularios de palabras.
Cada lote de entrenamiento consta de un conjunto de pares de oraciones que se agrupan con longitudes de secuencia similares. Cada lote de entrenamiento de pares de oraciones contiene alrededor de 25 000 tokens del idioma de origen y 25 000 tokens del idioma de destino.
3.4 Hardware y programación
Entrenamos a nuestros modelos en una máquina equipada con 8 GPU NVIDIA P100. Usando los modelos base con los hiperparámetros descritos en el documento, cada paso de entrenamiento toma alrededor de 0,4 segundos. Entrenamos los modelos base para un total de 100.000 pasos o 12 horas. Para modelos grandes, el tiempo de entrenamiento de cada paso es de 1,0 segundos, y los modelos grandes entrenaron 300.000 pasos (3,5 días).
3.5 Optimizador
Usamos el optimizador Adam, donde β 1 = 0.9 \beta_1=0.9b1=0,9 ,β2 = 0,98 \beta_2=0,98b2=0,98 yϵ = 1 0 − 9 ϵ=10^{-9}ϵ=1 0−9 . _ Variamos la tasa de aprendizaje durante el entrenamiento de acuerdo con la siguiente fórmula:
lrate = dmodel − 0.5 ⋅ min (paso _ num − 0.5 , stem _ num ⋅ calentamiento _ pasos − 1.5 ) lrate = d^{-0.5}_{model}\cdot min(step\_num^{-0.5}, stem \_num \cdot calentamiento\_pasos^{-1.5})tasa _ _ _ _=dmodelo _ _ _ _− 0,5⋅min ( paso _ n u m _ _ _− 0,5 ,tallo _ n u m _ _ _⋅Calentamiento _ pasos _ _ _ _ _ _ _ _ _− 1.5 )
Esto corresponde al primer warmup_steps warmup\_stepswarmup_steps aumenta la tasa de aprendizaje linealmente para s pasos y luego la disminuye proporcionalmente a la raíz cuadrada del número de pasos . Usamoswarmup_steps=4000 warmup\_steps=4000Calentamiento _ pasos _ _ _ _ _ _ _ _ _=4000 _
NOTA: Esta parte es muy importante. Esta configuración de modelo es necesaria para el entrenamiento.
Ejemplo de curvas del modelo para diferentes tamaños de modelo e hiperparámetros optimizados.
def rate(step, model_size, factor, warmup):
"""
we have to default the step to 1 for LambdaLR function
to avoid zero raising to negative power.
"""
if step == 0:
step = 1
return factor * (
model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
)
def example_learning_schedule():
opts = [
[512, 1, 4000], # example 1
[512, 1, 8000], # example 2
[256, 1, 4000], # example 3
]
dummy_model = torch.nn.Linear(1, 1)
learning_rates = []
# we have 3 examples in opts list.
for idx, example in enumerate(opts):
# run 20000 epoch for each example
optimizer = torch.optim.Adam(
dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9
)
lr_scheduler = LambdaLR(
optimizer=optimizer, lr_lambda=lambda step: rate(step, *example)
)
tmp = []
# take 20K dummy training steps, save the learning rate at each step
for step in range(20000):
tmp.append(optimizer.param_groups[0]["lr"])
optimizer.step()
lr_scheduler.step()
learning_rates.append(tmp)
learning_rates = torch.tensor(learning_rates)
# Enable altair to handle more than 5000 rows
alt.data_transformers.disable_max_rows()
opts_data = pd.concat(
[
pd.DataFrame(
{
"Learning Rate": learning_rates[warmup_idx, :],
"model_size:warmup": ["512:4000", "512:8000", "256:4000"][
warmup_idx
],
"step": range(20000),
}
)
for warmup_idx in [0, 1, 2]
]
)
return (
alt.Chart(opts_data)
.mark_line()
.properties(width=600)
.encode(x="step", y="Learning Rate", color="model_size:warmup:N")
.interactive()
)
example_learning_schedule()
3.6 Regularización
Suavizado de etiquetas
Durante el proceso de entrenamiento, el valor de suavizado de la etiqueta que usamos es ϵ ls = 0.1 \epsilon_{ls}=0.1ϵl s=0.1 (citar). Esto hace que el modelo sea menos comprensible ya que aprende a ser más incierto, pero mejora la precisión y las puntuaciones BLEU.
Usamos la pérdida div KL para el suavizado de etiquetas. En lugar de utilizar una distribución one-hot one-hot, creamos una distribución que tiene confianza en la palabra correcta y el resto de la masa suavizante distribuida por todo el vocabulario. Esta distribución tiene una "confianza" en la palabra correcta y el resto de la calidad de "suavizado" distribuido en todo el vocabulario.
class LabelSmoothing(nn.Module):
"Implement label smoothing."
def __init__(self, size, padding_idx, smoothing=0.0):
super(LabelSmoothing, self).__init__()
self.criterion = nn.KLDivLoss(reduction="sum")
self.padding_idx = padding_idx
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
self.size = size
self.true_dist = None
def forward(self, x, target):
assert x.size(1) == self.size
true_dist = x.data.clone()
true_dist.fill_(self.smoothing / (self.size - 2))
true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
true_dist[:, self.padding_idx] = 0
mask = torch.nonzero(target.data == self.padding_idx)
if mask.dim() > 0:
true_dist.index_fill_(0, mask.squeeze(), 0.0)
self.true_dist = true_dist
return self.criterion(x, true_dist.clone().detach())
Aquí podemos ver un ejemplo de cómo se distribuye la masa a las palabras en función de la confianza.
# Example of label smoothing.
def example_label_smoothing():
crit = LabelSmoothing(5, 0, 0.4)
predict = torch.FloatTensor(
[
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
]
)
crit(x=predict.log(), target=torch.LongTensor([2, 1, 0, 3, 3]))
LS_data = pd.concat(
[
pd.DataFrame(
{
"target distribution": crit.true_dist[x, y].flatten(),
"columns": y,
"rows": x,
}
)
for y in range(5)
for x in range(5)
]
)
return (
alt.Chart(LS_data)
.mark_rect(color="Blue", opacity=1)
.properties(height=200, width=200)
.encode(
alt.X("columns:O", title=None),
alt.Y("rows:O", title=None),
alt.Color(
"target distribution:Q", scale=alt.Scale(scheme="viridis")
),
)
.interactive()
)
show_example(example_label_smoothing)
El suavizado de etiquetas en realidad comienza a penalizar al modelo si tiene mucha confianza en una elección dada.
def loss(x, crit):
d = x + 3 * 1
predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d]])
return crit(predict.log(), torch.LongTensor([1])).data
def penalization_visualization():
crit = LabelSmoothing(5, 0, 0.1)
loss_data = pd.DataFrame(
{
"Loss": [loss(x, crit) for x in range(1, 100)],
"Steps": list(range(99)),
}
).astype("float")
return (
alt.Chart(loss_data)
.mark_line()
.properties(width=350)
.encode(
x="Steps",
y="Loss",
)
.interactive()
)
show_example(penalization_visualization)
4.Un primer ejemplo
Podemos comenzar probando una tarea de copia simple. Dado un conjunto de símbolos de entrada aleatorios de un vocabulario pequeño, el objetivo es generar estos mismos símbolos.
4.1 Datos sintéticos
def data_gen(V, batch_size, nbatches):
"Generate random data for a src-tgt copy task."
for i in range(nbatches):
data = torch.randint(1, V, size=(batch_size, 10))
data[:, 0] = 1
src = data.requires_grad_(False).clone().detach()
tgt = data.requires_grad_(False).clone().detach()
yield Batch(src, tgt, 0)
4.2 Cálculo de pérdidas
class SimpleLossCompute:
"A simple loss compute and train function."
def __init__(self, generator, criterion):
self.generator = generator
self.criterion = criterion
def __call__(self, x, y, norm):
x = self.generator(x)
sloss = (
self.criterion(
x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
)
/ norm
)
return sloss.data * norm, sloss
4.3 Decodificación codiciosa
Para simplificar, este código predice las traducciones usando decodificación codiciosa.
def greedy_decode(model, src, src_mask, max_len, start_symbol):
memory = model.encode(src, src_mask)
ys = torch.zeros(1, 1).fill_(start_symbol).type_as(src.data)
for i in range(max_len - 1):
out = model.decode(
memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
)
prob = model.generator(out[:, -1])
_, next_word = torch.max(prob, dim=1)
next_word = next_word.data[0]
ys = torch.cat(
[ys, torch.zeros(1, 1).type_as(src.data).fill_(next_word)], dim=1
)
return ys
# Train the simple copy task.
def example_simple_model():
V = 11
criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
model = make_model(V, V, N=2)
optimizer = torch.optim.Adam(
model.parameters(), lr=0.5, betas=(0.9, 0.98), eps=1e-9
)
lr_scheduler = LambdaLR(
optimizer=optimizer,
lr_lambda=lambda step: rate(
step, model_size=model.src_embed[0].d_model, factor=1.0, warmup=400
),
)
batch_size = 80
for epoch in range(20):
model.train()
run_epoch(
data_gen(V, batch_size, 20),
model,
SimpleLossCompute(model.generator, criterion),
optimizer,
lr_scheduler,
mode="train",
)
model.eval()
run_epoch(
data_gen(V, batch_size, 5),
model,
SimpleLossCompute(model.generator, criterion),
DummyOptimizer(),
DummyScheduler(),
mode="eval",
)[0]
model.eval()
src = torch.LongTensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
max_len = src.shape[1]
src_mask = torch.ones(1, 1, max_len)
print(greedy_decode(model, src, src_mask, max_len=max_len, start_symbol=0))
# execute_example(example_simple_model)
5.Un ejemplo del mundo real
Ahora consideramos un ejemplo del mundo real utilizando la tarea de traducción alemán-inglés de IWSLT. Esta tarea es mucho más pequeña que la tarea WMT considerada en el documento, pero esta tarea también puede ilustrar todo el sistema (de traducción). También mostramos cómo usar el procesamiento multi-GPU para que la tarea sea realmente rápida de entrenar.
5.1 Carga de datos
Usaremos torchtext y spacy para cargar el conjunto de datos para la tokenización.
# Load spacy tokenizer models, download them if they haven't been
# downloaded already
def load_tokenizers():
try:
spacy_de = spacy.load("de_core_news_sm")
except IOError:
os.system("python -m spacy download de_core_news_sm")
spacy_de = spacy.load("de_core_news_sm")
try:
spacy_en = spacy.load("en_core_web_sm")
except IOError:
os.system("python -m spacy download en_core_web_sm")
spacy_en = spacy.load("en_core_web_sm")
return spacy_de, spacy_en
def tokenize(text, tokenizer):
return [tok.text for tok in tokenizer.tokenizer(text)]
def yield_tokens(data_iter, tokenizer, index):
for from_to_tuple in data_iter:
yield tokenizer(from_to_tuple[index])
def build_vocabulary(spacy_de, spacy_en):
def tokenize_de(text):
return tokenize(text, spacy_de)
def tokenize_en(text):
return tokenize(text, spacy_en)
print("Building German Vocabulary ...")
train, val, test = datasets.Multi30k(language_pair=("de", "en"))
vocab_src = build_vocab_from_iterator(
yield_tokens(train + val + test, tokenize_de, index=0),
min_freq=2,
specials=["<s>", "</s>", "<blank>", "<unk>"],
)
print("Building English Vocabulary ...")
train, val, test = datasets.Multi30k(language_pair=("de", "en"))
vocab_tgt = build_vocab_from_iterator(
yield_tokens(train + val + test, tokenize_en, index=1),
min_freq=2,
specials=["<s>", "</s>", "<blank>", "<unk>"],
)
vocab_src.set_default_index(vocab_src["<unk>"])
vocab_tgt.set_default_index(vocab_tgt["<unk>"])
return vocab_src, vocab_tgt
def load_vocab(spacy_de, spacy_en):
if not exists("vocab.pt"):
vocab_src, vocab_tgt = build_vocabulary(spacy_de, spacy_en)
torch.save((vocab_src, vocab_tgt), "vocab.pt")
else:
vocab_src, vocab_tgt = torch.load("vocab.pt")
print("Finished.\nVocabulary sizes:")
print(len(vocab_src))
print(len(vocab_tgt))
return vocab_src, vocab_tgt
if is_interactive_notebook():
# global variables used later in the script
spacy_de, spacy_en = show_example(load_tokenizers)
vocab_src, vocab_tgt = show_example(load_vocab, args=[spacy_de, spacy_en])
Finished.
Vocabulary sizes:
59981
36745
La dosificación es muy importante para la velocidad de entrenamiento. Queremos tener lotes muy parejos con un relleno mínimo absoluto. Para hacer esto, tenemos que hacer algunas modificaciones al procesamiento por lotes predeterminado de torchtext. Este código parchea el procesamiento por lotes predeterminado de torchtext para garantizar que encontremos lotes estables al buscar suficientes oraciones.
5.2 Iteradores
def collate_batch(
batch,
src_pipeline,
tgt_pipeline,
src_vocab,
tgt_vocab,
device,
max_padding=128,
pad_id=2,
):
bs_id = torch.tensor([0], device=device) # <s> token id
eos_id = torch.tensor([1], device=device) # </s> token id
src_list, tgt_list = [], []
for (_src, _tgt) in batch:
processed_src = torch.cat(
[
bs_id,
torch.tensor(
src_vocab(src_pipeline(_src)),
dtype=torch.int64,
device=device,
),
eos_id,
],
0,
)
processed_tgt = torch.cat(
[
bs_id,
torch.tensor(
tgt_vocab(tgt_pipeline(_tgt)),
dtype=torch.int64,
device=device,
),
eos_id,
],
0,
)
src_list.append(
# warning - overwrites values for negative values of padding - len
pad(
processed_src,
(
0,
max_padding - len(processed_src),
),
value=pad_id,
)
)
tgt_list.append(
pad(
processed_tgt,
(0, max_padding - len(processed_tgt)),
value=pad_id,
)
)
src = torch.stack(src_list)
tgt = torch.stack(tgt_list)
return (src, tgt)
def create_dataloaders(
device,
vocab_src,
vocab_tgt,
spacy_de,
spacy_en,
batch_size=12000,
max_padding=128,
is_distributed=True,
):
# def create_dataloaders(batch_size=12000):
def tokenize_de(text):
return tokenize(text, spacy_de)
def tokenize_en(text):
return tokenize(text, spacy_en)
def collate_fn(batch):
return collate_batch(
batch,
tokenize_de,
tokenize_en,
vocab_src,
vocab_tgt,
device,
max_padding=max_padding,
pad_id=vocab_src.get_stoi()["<blank>"],
)
train_iter, valid_iter, test_iter = datasets.Multi30k(
language_pair=("de", "en")
)
train_iter_map = to_map_style_dataset(
train_iter
) # DistributedSampler needs a dataset len()
train_sampler = (
DistributedSampler(train_iter_map) if is_distributed else None
)
valid_iter_map = to_map_style_dataset(valid_iter)
valid_sampler = (
DistributedSampler(valid_iter_map) if is_distributed else None
)
train_dataloader = DataLoader(
train_iter_map,
batch_size=batch_size,
shuffle=(train_sampler is None),
sampler=train_sampler,
collate_fn=collate_fn,
)
valid_dataloader = DataLoader(
valid_iter_map,
batch_size=batch_size,
shuffle=(valid_sampler is None),
sampler=valid_sampler,
collate_fn=collate_fn,
)
return train_dataloader, valid_dataloader
5.3 Entrenamiento del sistema
def train_worker(
gpu,
ngpus_per_node,
vocab_src,
vocab_tgt,
spacy_de,
spacy_en,
config,
is_distributed=False,
):
print(f"Train worker process using GPU: {gpu} for training", flush=True)
torch.cuda.set_device(gpu)
pad_idx = vocab_tgt["<blank>"]
d_model = 512
model = make_model(len(vocab_src), len(vocab_tgt), N=6)
model.cuda(gpu)
module = model
is_main_process = True
if is_distributed:
dist.init_process_group(
"nccl", init_method="env://", rank=gpu, world_size=ngpus_per_node
)
model = DDP(model, device_ids=[gpu])
module = model.module
is_main_process = gpu == 0
criterion = LabelSmoothing(
size=len(vocab_tgt), padding_idx=pad_idx, smoothing=0.1
)
criterion.cuda(gpu)
train_dataloader, valid_dataloader = create_dataloaders(
gpu,
vocab_src,
vocab_tgt,
spacy_de,
spacy_en,
batch_size=config["batch_size"] // ngpus_per_node,
max_padding=config["max_padding"],
is_distributed=is_distributed,
)
optimizer = torch.optim.Adam(
model.parameters(), lr=config["base_lr"], betas=(0.9, 0.98), eps=1e-9
)
lr_scheduler = LambdaLR(
optimizer=optimizer,
lr_lambda=lambda step: rate(
step, d_model, factor=1, warmup=config["warmup"]
),
)
train_state = TrainState()
for epoch in range(config["num_epochs"]):
if is_distributed:
train_dataloader.sampler.set_epoch(epoch)
valid_dataloader.sampler.set_epoch(epoch)
model.train()
print(f"[GPU{gpu}] Epoch {epoch} Training ====", flush=True)
_, train_state = run_epoch(
(Batch(b[0], b[1], pad_idx) for b in train_dataloader),
model,
SimpleLossCompute(module.generator, criterion),
optimizer,
lr_scheduler,
mode="train+log",
accum_iter=config["accum_iter"],
train_state=train_state,
)
GPUtil.showUtilization()
if is_main_process:
file_path = "%s%.2d.pt" % (config["file_prefix"], epoch)
torch.save(module.state_dict(), file_path)
torch.cuda.empty_cache()
print(f"[GPU{gpu}] Epoch {epoch} Validation ====", flush=True)
model.eval()
sloss = run_epoch(
(Batch(b[0], b[1], pad_idx) for b in valid_dataloader),
model,
SimpleLossCompute(module.generator, criterion),
DummyOptimizer(),
DummyScheduler(),
mode="eval",
)
print(sloss)
torch.cuda.empty_cache()
if is_main_process:
file_path = "%sfinal.pt" % config["file_prefix"]
torch.save(module.state_dict(), file_path)
def train_distributed_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config):
from the_annotated_transformer import train_worker
ngpus = torch.cuda.device_count()
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12356"
print(f"Number of GPUs detected: {ngpus}")
print("Spawning training processes ...")
mp.spawn(
train_worker,
nprocs=ngpus,
args=(ngpus, vocab_src, vocab_tgt, spacy_de, spacy_en, config, True),
)
def train_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config):
if config["distributed"]:
train_distributed_model(
vocab_src, vocab_tgt, spacy_de, spacy_en, config
)
else:
train_worker(
0, 1, vocab_src, vocab_tgt, spacy_de, spacy_en, config, False
)
def load_trained_model():
config = {
"batch_size": 32,
"distributed": False,
"num_epochs": 8,
"accum_iter": 10,
"base_lr": 1.0,
"max_padding": 72,
"warmup": 3000,
"file_prefix": "multi30k_model_",
}
model_path = "multi30k_model_final.pt"
if not exists(model_path):
train_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config)
model = make_model(len(vocab_src), len(vocab_tgt), N=6)
model.load_state_dict(torch.load("multi30k_model_final.pt"))
return model
if is_interactive_notebook():
model = load_trained_model()
Una vez entrenado, podemos decodificar el modelo para generar un conjunto de traducciones. Aquí simplemente traducimos la primera oración en el conjunto de validación. Este conjunto de datos es muy pequeño, por lo que los resultados de la traducción de la búsqueda codiciosa son bastante precisos.
6. Componentes adicionales: BPE, búsqueda, promedio
El contenido anterior cubre principalmente el modelo de transformador en sí, pero en realidad hay cuatro funciones adicionales que no hemos cubierto. Pero implementamos todas estas funciones adicionales en OpenNMT-py .
- BPE/Word-piece : Podemos usar una biblioteca para preprocesar primero los datos en unidades de subpalabras. Consulte la subpalabra-nmt de Rico Sennrich para la implementación. Estos modelos transforman los datos de entrenamiento de la siguiente manera:
▁Die ▁Protokoll datei ▁kann ▁ heimlich ▁per ▁E - Mail ▁oder ▁FTP ▁an ▁einen ▁bestimmte n ▁Empfänger ▁gesendet ▁werden - Incrustaciones compartidas : cuando usamos BPE con vocabulario compartido, podemos compartir el mismo vector de peso entre fuente/objetivo/generador. Ver (citas) para más detalles. Para agregarlo a su modelo, simplemente haga lo siguiente:
if False:
model.src_embed[0].lut.weight = model.tgt_embeddings[0].lut.weight
model.generator.lut.weight = model.tgt_embed[0].lut.weight
- Beam Search: Esto es un poco demasiado complejo para cubrirlo aquí. Ver OpenNMT-py para una implementación de pytorch .
- Promedio del modelo: el documento promedia los últimos k puntos de control para obtener el efecto de integración. Si tenemos un montón de modelos, podemos hacer esto:
def average(model, models):
"Average models into model"
for ps in zip(*[m.params() for m in [model] + models]):
ps[0].copy_(torch.sum(*ps[1:]) / len(ps[1:]))
7.Resultados
En la tarea de traducción inglés-alemán de WMT 2014, el modelo de transformador grande (Transformador (grande) en la Tabla 2) supera a los mejores modelos informados anteriormente (incluidos los modelos de conjuntos) en más de 2,0 BLEU, y la nueva puntuación BLEU más alta es 28,4. La configuración del modelo se muestra en la parte inferior de la Tabla 3, y el modelo se entrenó durante 3,5 días en 8 GPU P100. Incluso nuestro modelo base supera a los modelos anteriores y a los modelos de conjuntos con un costo de capacitación mucho menor que los modelos anteriores.
En la tarea de traducción inglés-francés de WMT 2014, nuestro modelo grande logra una puntuación BLEU de 41,0 puntos, superando a todos los modelos individuales lanzados anteriormente, y el costo de capacitación es menos de 1/4 del modelo de última generación anterior. El modelo Transformer (grande) entrenado en la tarea de traducción inglés-francés utiliza una tasa de abandono de Pdrop = 0,1 en lugar de 0,3.
Con varias funciones adicionales mencionadas en la sección anterior, OpenNMT-py puede alcanzar una puntuación de 26,9 en EN-DE WMT. A continuación, cargo estos parámetros en mi código para reproducirlos.
# Load data and model for output checks
def check_outputs(
valid_dataloader,
model,
vocab_src,
vocab_tgt,
n_examples=15,
pad_idx=2,
eos_string="</s>",
):
results = [()] * n_examples
for idx in range(n_examples):
print("\nExample %d ========\n" % idx)
b = next(iter(valid_dataloader))
rb = Batch(b[0], b[1], pad_idx)
greedy_decode(model, rb.src, rb.src_mask, 64, 0)[0]
src_tokens = [
vocab_src.get_itos()[x] for x in rb.src[0] if x != pad_idx
]
tgt_tokens = [
vocab_tgt.get_itos()[x] for x in rb.tgt[0] if x != pad_idx
]
print(
"Source Text (Input) : "
+ " ".join(src_tokens).replace("\n", "")
)
print(
"Target Text (Ground Truth) : "
+ " ".join(tgt_tokens).replace("\n", "")
)
model_out = greedy_decode(model, rb.src, rb.src_mask, 72, 0)[0]
model_txt = (
" ".join(
[vocab_tgt.get_itos()[x] for x in model_out if x != pad_idx]
).split(eos_string, 1)[0]
+ eos_string
)
print("Model Output : " + model_txt.replace("\n", ""))
results[idx] = (rb, src_tokens, tgt_tokens, model_out, model_txt)
return results
def run_model_example(n_examples=5):
global vocab_src, vocab_tgt, spacy_de, spacy_en
print("Preparing Data ...")
_, valid_dataloader = create_dataloaders(
torch.device("cpu"),
vocab_src,
vocab_tgt,
spacy_de,
spacy_en,
batch_size=1,
is_distributed=False,
)
print("Loading Trained Model ...")
model = make_model(len(vocab_src), len(vocab_tgt), N=6)
model.load_state_dict(
torch.load("multi30k_model_final.pt", map_location=torch.device("cpu"))
)
print("Checking Model Outputs:")
example_data = check_outputs(
valid_dataloader, model, vocab_src, vocab_tgt, n_examples=n_examples
)
return model, example_data
# execute_example(run_model_example)
7.1 Visualización de atención
Incluso con una decodificación codiciosa, la traducción se ve bien. Podemos visualizar esto más a fondo para ver qué está pasando en cada nivel de atención.
def mtx2df(m, max_row, max_col, row_tokens, col_tokens):
"convert a dense matrix to a data frame with row and column indices"
return pd.DataFrame(
[
(
r,
c,
float(m[r, c]),
"%.3d %s"
% (r, row_tokens[r] if len(row_tokens) > r else "<blank>"),
"%.3d %s"
% (c, col_tokens[c] if len(col_tokens) > c else "<blank>"),
)
for r in range(m.shape[0])
for c in range(m.shape[1])
if r < max_row and c < max_col
],
# if float(m[r,c]) != 0 and r < max_row and c < max_col],
columns=["row", "column", "value", "row_token", "col_token"],
)
def attn_map(attn, layer, head, row_tokens, col_tokens, max_dim=30):
df = mtx2df(
attn[0, head].data,
max_dim,
max_dim,
row_tokens,
col_tokens,
)
return (
alt.Chart(data=df)
.mark_rect()
.encode(
x=alt.X("col_token", axis=alt.Axis(title="")),
y=alt.Y("row_token", axis=alt.Axis(title="")),
color="value",
tooltip=["row", "column", "value", "row_token", "col_token"],
)
.properties(height=400, width=400)
.interactive()
)
def get_encoder(model, layer):
return model.encoder.layers[layer].self_attn.attn
def get_decoder_self(model, layer):
return model.decoder.layers[layer].self_attn.attn
def get_decoder_src(model, layer):
return model.decoder.layers[layer].src_attn.attn
def visualize_layer(model, layer, getter_fn, ntokens, row_tokens, col_tokens):
# ntokens = last_example[0].ntokens
attn = getter_fn(model, layer)
n_heads = attn.shape[1]
charts = [
attn_map(
attn,
0,
h,
row_tokens=row_tokens,
col_tokens=col_tokens,
max_dim=ntokens,
)
for h in range(n_heads)
]
assert n_heads == 8
return alt.vconcat(
charts[0]
# | charts[1]
| charts[2]
# | charts[3]
| charts[4]
# | charts[5]
| charts[6]
# | charts[7]
# layer + 1 due to 0-indexing
).properties(title="Layer %d" % (layer + 1))
7.2 Atención propia del codificador
def viz_encoder_self():
model, example_data = run_model_example(n_examples=1)
example = example_data[
len(example_data) - 1
] # batch object for the final example
layer_viz = [
visualize_layer(
model, layer, get_encoder, len(example[1]), example[1], example[1]
)
for layer in range(6)
]
return alt.hconcat(
layer_viz[0]
# & layer_viz[1]
& layer_viz[2]
# & layer_viz[3]
& layer_viz[4]
# & layer_viz[5]
)
show_example(viz_encoder_self)
Preparing Data ...
Loading Trained Model ...
Checking Model Outputs:
Example 0 ========
Source Text (Input) : <s> Zwei Frauen in pinkfarbenen T-Shirts und <unk> unterhalten sich vor einem <unk> . </s>
Target Text (Ground Truth) : <s> Two women wearing pink T - shirts and blue jeans converse outside clothing store . </s>
Model Output : <s> Two women in pink shirts and face are talking in front of a <unk> . </s>
7.3 Atención propia del decodificador
def viz_decoder_self():
model, example_data = run_model_example(n_examples=1)
example = example_data[len(example_data) - 1]
layer_viz = [
visualize_layer(
model,
layer,
get_decoder_self,
len(example[1]),
example[1],
example[1],
)
for layer in range(6)
]
return alt.hconcat(
layer_viz[0]
& layer_viz[1]
& layer_viz[2]
& layer_viz[3]
& layer_viz[4]
& layer_viz[5]
)
show_example(viz_decoder_self)
Preparing Data ...
Loading Trained Model ...
Checking Model Outputs:
Example 0 ========
Source Text (Input) : <s> Eine Gruppe von Männern in Kostümen spielt Musik . </s>
Target Text (Ground Truth) : <s> A group of men in costume play music . </s>
Model Output : <s> A group of men in costumes playing music . </s>
7.4 Atención de origen del decodificador
def viz_decoder_src():
model, example_data = run_model_example(n_examples=1)
example = example_data[len(example_data) - 1]
layer_viz = [
visualize_layer(
model,
layer,
get_decoder_src,
max(len(example[1]), len(example[2])),
example[1],
example[2],
)
for layer in range(6)
]
return alt.hconcat(
layer_viz[0]
& layer_viz[1]
& layer_viz[2]
& layer_viz[3]
& layer_viz[4]
& layer_viz[5]
)
show_example(viz_decoder_src)
Preparing Data ...
Loading Trained Model ...
Checking Model Outputs:
Example 0 ========
Source Text (Input) : <s> Ein kleiner Junge verwendet einen Bohrer , um ein Loch in ein Holzstück zu machen . </s>
Target Text (Ground Truth) : <s> A little boy using a drill to make a hole in a piece of wood . </s>
Model Output : <s> A little boy uses a machine to be working in a hole in a log . </s>
8.Conclusión
Espero que este código sea útil para futuras investigaciones. Si tiene alguna pregunta, por favor contáctenos.