def TransformerEncoder(vocab_size,
n_classes=10,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Returns a Transformer encoder model.
The input to the model is a tensor of tokens.
Args:
vocab_size: int: vocab size
n_classes: how many classes on output
d_model: int: depth of embedding
d_ff: int: depth of feed-forward layer
n_layers: int: number of encoder/decoder layers
n_heads: int: number of attention heads
dropout: float: dropout rate (how much to drop out)
max_len: int: maximum symbol length for positional encoding
mode: str: 'train' or 'eval'
Returns:
A Transformer model as a layer that maps from a tensor of tokens to
activations over a set of output classes.
"""
embedder = [
tl.Embedding(d_model, vocab_size),
tl.Dropout(rate=dropout, name='emb_dropout', mode=mode),
tl.PositionalEncoding(max_len=max_len),
]
return tl.Model([ # tokens
tl.Dup(), # toks toks
tl.Parallel(embedder, tl.PaddingMask()), # vecs mask
[
EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode)
for i in range(n_layers)
], # vecs mask
tl.Parallel([], tl.Drop()), # ____ 0
tl.LayerNorm(), # vecs
tl.Mean(axis=1), # Average on length. # vecs
tl.Dense(n_classes), # vecs
tl.LogSoftmax(), # vecs
])
def TransformerEncoder(vocab_size,
n_classes=10,
d_feature=512,
d_feedforward=2048,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Transformer encoder.
Args:
vocab_size: int: vocab size
n_classes: how many classes on output
d_feature: int: depth of embedding
d_feedforward: int: depth of feed-forward layer
n_layers: int: number of encoder/decoder layers
n_heads: int: number of attention heads
dropout: float: dropout rate (how much to drop out)
max_len: int: maximum symbol length for positional encoding
mode: str: 'train' or 'eval'
Returns:
the Transformer encoder layer.
"""
positional_embedder = [
tl.Embedding(d_feature, vocab_size),
tl.Dropout(rate=dropout, mode=mode),
tl.PositionalEncoding(max_len=max_len),
]
return [
tl.Branch(positional_embedder, tl.PaddingMask()), # Create mask.
[
EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
for _ in range(n_layers)
],
tl.Select(0), # Drop mask.
tl.LayerNorm(),
tl.Mean(axis=1), # Average on length.
tl.Dense(n_classes),
tl.LogSoftmax(),
]
def TransformerEncoder(vocab_size,
num_classes=10,
feature_depth=512,
feedforward_depth=2048,
num_layers=6,
num_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Transformer encoder.
Args:
vocab_size: int: vocab size
num_classes: how many classes on output
feature_depth: int: depth of embedding
feedforward_depth: int: depth of feed-forward layer
num_layers: int: number of encoder/decoder layers
num_heads: int: number of attention heads
dropout: float: dropout rate (how much to drop out)
max_len: int: maximum symbol length for positional encoding
mode: str: 'train' or 'eval'
Returns:
the Transformer encoder layer.
"""
input_embedding = layers.Serial(
layers.Embedding(feature_depth, vocab_size),
layers.Dropout(rate=dropout, mode=mode),
layers.PositionalEncoding(max_len=max_len)
)
return layers.Serial(
layers.Branch(), # Branch input to create embedding and mask.
layers.Parallel(input_embedding, layers.PaddingMask()),
layers.Serial(*[EncoderLayer(feature_depth, feedforward_depth, num_heads,
dropout, mode)
for _ in range(num_layers)]),
layers.FirstBranch(), # Drop the mask.
layers.LayerNorm(),
layers.Mean(axis=1), # Average on length.
layers.Dense(num_classes),
layers.LogSoftmax()
)