def _TransformerParamsWithEmbeddings(self,
num_decoder_layers=0,
num_encoder_layers=4,
splits=1,
num_micro_batches=1):
model_dim = 4
params = GPipeTransformerStack.Params()
params.name = 'transformer'
params.model_dim = model_dim
params.num_decoder_layers = num_decoder_layers
params.decoder_tpl.tr_atten_tpl.num_attention_heads = 1
params.decoder_tpl.tr_fflayer_tpl.hidden_dim = model_dim
params.num_encoder_layers = num_encoder_layers
params.encoder_tpl.tr_atten_tpl.num_attention_heads = 1
params.encoder_tpl.tr_fflayer_tpl.hidden_dim = model_dim
params.num_micro_batches = num_micro_batches
params.use_pipelined_embeddings = True
params.state_dtype = tf.float32
emb_params = params.emb_tpl
# Default config for the token embedding.
emb_params.token_emb.use_matmul = True
emb_params.token_emb.use_3d_weight_tensor = False
emb_params.token_emb.vocab_size = 10
emb_params.token_emb.embedding_dim = model_dim
# Default config for the position embedding.
emb_params.position_emb.embedding_dim = model_dim
emb_params.position_emb.trainable_scaling = False
params.splits = splits
params.random_seed = 0
return params
def _TransformerParamsWithEmbeddings(num_decoder_layers=0,
num_encoder_layers=4,
splits=1,
num_micro_batches=1,
has_softmax=False,
use_task_ids=False):
model_dim = 4
params = GPipeTransformerStack.Params()
params.name = 'transformer'
params.model_dim = model_dim
params.num_decoder_layers = num_decoder_layers
params.decoder_tpl.source_dim = model_dim
params.decoder_tpl.tr_atten_tpl.num_attention_heads = 1
params.decoder_tpl.tr_fflayer_tpl.hidden_dim = model_dim
params.num_encoder_layers = num_encoder_layers
params.encoder_tpl.source_dim = model_dim
params.encoder_tpl.tr_atten_tpl.num_attention_heads = 1
params.encoder_tpl.tr_fflayer_tpl.hidden_dim = model_dim
params.num_micro_batches = num_micro_batches
params.state_dtype = tf.float32
if has_softmax:
params.softmax_tpl.input_dim = model_dim
params.softmax_tpl.num_classes = 2
else:
params.softmax_tpl = None
emb_params = params.emb_tpl
# Default config for the token embedding.
emb_params.token_emb.use_matmul = True
emb_params.token_emb.use_3d_weight_tensor = False
emb_params.token_emb.vocab_size = 10
emb_params.token_emb.embedding_dim = model_dim
# Default config for the position embedding.
emb_params.position_emb.embedding_dim = model_dim
emb_params.position_emb.trainable_scaling = False
# Task embeddings.
if use_task_ids:
emb_params.enc_task_emb = emb_params.token_emb.Copy()
emb_params.dec_task_emb = emb_params.token_emb.Copy()
params.splits = splits
params.random_seed = 0
return params
def _TransformerParams(self,
num_decoder_layers=0,
num_encoder_layers=4,
splits=1,
num_micro_batches=1):
model_dim = 2
params = GPipeTransformerStack.Params()
params.name = 'transformer'
params.model_dim = model_dim
params.num_decoder_layers = num_decoder_layers
params.decoder_tpl.tr_atten_tpl.num_attention_heads = 1
params.decoder_tpl.tr_fflayer_tpl.hidden_dim = model_dim
params.num_encoder_layers = num_encoder_layers
params.encoder_tpl.tr_atten_tpl.num_attention_heads = 1
params.encoder_tpl.tr_fflayer_tpl.hidden_dim = model_dim
params.num_micro_batches = num_micro_batches
params.splits = splits
params.random_seed = 0
return params