def __init__(
self,
n_heads,
n_layers,
embedding_size,
ffn_size,
vocabulary_size,
embedding=None,
dropout=0.0,
attention_dropout=0.0,
relu_dropout=0.0,
padding_idx=0,
learn_positional_embeddings=False,
embeddings_scale=False,
reduction_type='mean',
n_positions=1024,
activation='relu',
variant='aiayn',
n_segments=0,
output_scaling=1.0,
):
super(TransformerEncoder, self).__init__()
self.embedding_size = embedding_size
self.ffn_size = ffn_size
self.n_layers = n_layers
self.n_heads = n_heads
self.dim = embedding_size
self.embeddings_scale = embeddings_scale
self.reduction_type = reduction_type
self.padding_idx = padding_idx
# this is --dropout, not --relu-dropout or --attention-dropout
self.dropout = nn.Dropout(p=dropout)
self.variant = variant
self.n_segments = n_segments
self.n_positions = n_positions
self.out_dim = embedding_size
assert (embedding_size % n_heads == 0
), 'Transformer embedding size must be a multiple of n_heads'
# check input formats:
if embedding is not None:
assert (embedding_size is None
or embedding_size == embedding.weight.shape[1]
), "Embedding dim must match the embedding size."
if embedding is not None:
self.embeddings = embedding
else:
raise AssertionError(
"This code should not execute. Left here in case we want to enable it."
)
assert padding_idx is not None
self.embeddings = nn.Embedding(vocabulary_size,
embedding_size,
padding_idx=padding_idx)
nn.init.normal_(self.embeddings.weight, 0, embedding_size**-0.5)
# create the positional embeddings
self.position_embeddings = nn.Embedding(n_positions, embedding_size)
if not learn_positional_embeddings:
create_position_codes(n_positions,
embedding_size,
out=self.position_embeddings.weight)
else:
nn.init.normal_(self.position_embeddings.weight, 0,
embedding_size**-0.5)
# embedding normalization
if self.variant == 'xlm':
self.norm_embeddings = LayerNorm(self.dim, eps=LAYER_NORM_EPS)
elif self.variant == 'aiayn':
pass
else:
raise ValueError("Can't handle --variant {}".format(self.variant))
if self.n_segments >= 1:
self.segment_embeddings = nn.Embedding(self.n_segments, self.dim)
# build the model
self.layers = nn.ModuleList()
for _ in range(self.n_layers):
self.layers.append(
TransformerEncoderLayer(
n_heads,
embedding_size,
ffn_size,
attention_dropout=attention_dropout,
relu_dropout=relu_dropout,
dropout=dropout,
variant=variant,
activation=activation,
))
self.output_scaling = output_scaling
def __init__(self,
num_layers,
vocab_size,
hidden_size,
num_attention_heads,
embedding_dropout_prob,
attention_dropout_prob,
output_dropout_prob,
layernorm_epsilon,
max_sequence_length,
checkpoint_activations,
checkpoint_num_layers=1,
parallel_output=True,
num_experts=1,
type_vocab_size=2):
super(BertMixtureModel_v0, self).__init__()
self.parallel_output = parallel_output
init_method = init_method_normal(std=0.02)
# Word embeddings (parallel).
self.word_embeddings = mpu.VocabParallelEmbedding(
vocab_size, hidden_size, init_method=init_method)
# Position embedding (serial).
self.position_embeddings = torch.nn.Embedding(max_sequence_length,
hidden_size)
# Initialize the position embeddings.
init_method(self.position_embeddings.weight)
# Token Type Enbeddings.
self.token_type_embeddings = torch.nn.Embedding(
type_vocab_size, hidden_size)
# Initialize the token type embeddings.
init_method(self.token_type_embeddings.weight)
# Embeddings dropout
self.embedding_dropout = torch.nn.Dropout(embedding_dropout_prob)
self.input_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon)
self.hrs_head = torch.nn.Linear(hidden_size, 1)
init_method(self.hrs_head.weight)
self.click_head = torch.nn.Linear(hidden_size, 1)
init_method(self.click_head.weight)
# Transformer
self.transformer = mpu.BertParallelTransformer(num_layers,
hidden_size,
num_attention_heads,
attention_dropout_prob,
output_dropout_prob,
checkpoint_activations,
checkpoint_num_layers,
num_experts=num_experts)
self.dropout = torch.nn.Dropout(output_dropout_prob)
# limitations under the License.
# Parts of the code here are adapted from PyTorch
# repo: https://github.com/pytorch/pytorch
import math
import torch
import torch.nn.functional as F
import torch.nn.init as init
from torch.nn.parameter import Parameter
try:
from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm
# Try to use FusedLayerNorm from Apex - this will trigger an error.
_ = LayerNorm(8, eps=1e-5)
except Exception as e:
print('WARNING: APEX is not installed, using torch.nn.LayerNorm '
'instead of apex.normalization.FusedLayerNorm!')
from torch.nn import LayerNorm
from .initialize import get_model_parallel_rank
from .initialize import get_model_parallel_world_size
from .mappings import copy_to_model_parallel_region
from .mappings import gather_from_model_parallel_region
from .mappings import reduce_from_model_parallel_region
from .mappings import scatter_to_model_parallel_region
from .random import get_cuda_rng_tracker
from .utils import divide
from .utils import split_tensor_along_last_dim
def __init__(
self,
n_heads,
n_layers,
embedding_size,
ffn_size,
vocabulary_size,
embedding=None,
dropout=0.0,
attention_dropout=0.0,
relu_dropout=0.0,
embeddings_scale=True,
learn_positional_embeddings=False,
padding_idx=None,
n_positions=1024,
n_segments=0,
variant='aiayn',
activation='relu',
):
super().__init__()
self.embedding_size = embedding_size
self.ffn_size = ffn_size
self.n_layers = n_layers
self.n_heads = n_heads
self.dim = embedding_size
self.activation = activation
self.variant = variant
self.embeddings_scale = embeddings_scale
self.dropout = nn.Dropout(p=dropout) # --dropout
self.n_positions = n_positions
self.out_dim = embedding_size
assert (embedding_size % n_heads == 0
), 'Transformer embedding size must be a multiple of n_heads'
self.embeddings = embedding
if self.variant == 'xlm' or self.variant == 'prelayernorm':
self.norm_embeddings = LayerNorm(self.dim, eps=LAYER_NORM_EPS)
elif self.variant == 'aiayn':
pass
else:
raise ValueError("Can't handle --variant {}".format(self.variant))
# create the positional embeddings
self.position_embeddings = nn.Embedding(n_positions, embedding_size)
if not learn_positional_embeddings:
create_position_codes(n_positions,
embedding_size,
out=self.position_embeddings.weight)
else:
nn.init.normal_(self.position_embeddings.weight, 0,
embedding_size**-0.5)
# build the model
self.layers = nn.ModuleList()
for _ in range(self.n_layers):
self.layers.append(
TransformerDecoderLayer(
n_heads,
embedding_size,
ffn_size,
attention_dropout=attention_dropout,
relu_dropout=relu_dropout,
dropout=dropout,
activation=activation,
variant=variant,
))
def __init__(self, config):
super(TransformerOutput, self).__init__()
self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
self.LayerNorm = LayerNorm(config.hidden_size,
eps=config.layer_norm_eps)
self.dropout = nn.Dropout(config.hidden_dropout_prob)