np_datatype = np.float32
if args.data_type == "fp16":
tf_datatype = tf.float16
np_datatype = np.float16
use_XLA = args.use_XLA
beam_search_diversity_rate = args.beam_search_diversity_rate
sampling_topk = args.sampling_topk
sampling_topp = args.sampling_topp
hidden_dim = head_num * size_per_head
memory_hidden_dim = args.memory_hidden_dim
decoder_args = TransformerArgument(
beam_width=beam_width,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
kernel_init_range=kernel_initializer_range,
bias_init_range=bias_initializer_range)
decoding_args = DecodingBeamsearchArgument(vocab_size,
start_of_sentence_id,
end_of_sentence_id, max_seq_len,
decoder_args,
beam_search_diversity_rate)
decoder_args_2 = copy.deepcopy(decoder_args) # for beam search
decoder_args_2.__dict__ = copy.deepcopy(decoder_args.__dict__)
decoder_args_2.beam_width = 1 # for sampling
decoding_sampling_args = DecodingSamplingArgument(
if avg_seq_len != -1 and remove_padding == True:
# This means we use "remove_padding" and set a smaller average sequence length
sequence_length = np.ones(batch_size) * avg_seq_len
from_data = np.random.randn(batch_size, max_seq_len, hidden_dim)
from_tensor = tf.convert_to_tensor(from_data, dtype=tf_datatype)
attention_mask = build_sequence_mask(sequence_length,
num_heads=head_num,
maximum_length=max_seq_len,
dtype=tf_datatype)
encoder_notInt8_args = TransformerArgument(beam_width=1,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
remove_padding=remove_padding,
int8_mode=0)
encoder_Int8_v1_args = TransformerArgument(beam_width=1,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
remove_padding=remove_padding,
int8_mode=1)
encoder_Int8_v2_args = TransformerArgument(beam_width=1,
head_num=head_num,
size_per_head=size_per_head,
def sample_model(vocab_file="models/gpt2-vocab.json",
bpe_file="models/gpt2-merges.txt",
model_name='124M',
nsamples=1,
batch_size=1,
length=12,
temperature=1,
top_k=4,
top_p=0,
models_dir='models',
data_type='fp32'):
"""Run the sample_model.
:model_name=124M : String, which model to use
:nsamples=0 : Number of samples to return, if 0, continues to
generate samples indefinately.
:batch_size=1 : Number of batches (only affects speed/memory).
:length=None : Number of tokens in generated text, if None (default), is
determined by model hyperparameters
:temperature=1 : Float value controlling randomness in boltzmann
distribution. Lower temperature results in less random completions. As the
temperature approaches zero, the model will become deterministic and
repetitive. Higher temperature results in more random completions.
:top_k=4 : Integer value controlling diversity. 1 means only 1 word is
considered for each step (token), resulting in deterministic completions,
while 40 means 40 words are considered at each step. 0 (default) is a
special setting meaning no restrictions. 40 generally is a good value.
:models_dir : path to parent folder containing model subfolders
(i.e. contains the <model_name> folder)
"""
np.random.seed(1)
tf.set_random_seed(1)
if data_type == 'fp32':
tf_data_type = tf.float32
elif data_type == 'fp16':
tf_data_type = tf.float16
else:
assert (False)
models_dir = os.path.expanduser(os.path.expandvars(models_dir))
vocab_file = os.path.join(models_dir, model_name, 'encoder.json')
bpe_file = os.path.join(models_dir, model_name, 'vocab.bpe')
enc = encoder.get_encoder(vocab_file, bpe_file)
hparams = HParams(n_vocab=0, n_ctx=1024, n_embd=768, n_head=12, n_layer=12)
with open(os.path.join(models_dir, model_name, 'hparams.json')) as f:
hparams.override_from_dict(json.load(f))
if length is None:
length = hparams.n_ctx
elif length > hparams.n_ctx:
raise ValueError("Can't get samples longer than window size: %s" %
hparams.n_ctx)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(graph=tf.Graph(), config=config) as sess:
saver = tf.train.import_meta_graph("{}/{}/model.ckpt.meta".format(
models_dir, model_name))
lengths = np.random.randint(low=1, high=8, size=batch_size)
min_start_length = lengths.min()
max_start_length = lengths.max()
attention_mask = np.tile(np.tri(min_start_length), (batch_size, 1, 1))
start_ids = np.ones([batch_size, max_start_length
]) * enc.encoder['<|endoftext|>']
for i in range(batch_size):
start_ids[i][0:lengths[i]] = 198
# User can put some real start ids here, we use '\n' (198) here.
sess.run(tf.global_variables_initializer())
print("[INFO] restore the model {}/{}".format(models_dir, model_name))
saver.restore(sess,
("{}/{}/model.ckpt".format(models_dir, model_name)))
decoder_args = TransformerArgument(beam_width=1,
head_num=hparams.n_head,
size_per_head=hparams.n_embd //
hparams.n_head,
num_layer=hparams.n_layer,
dtype=tf_data_type,
kernel_init_range=0.00,
bias_init_range=0.00)
decoding_args = DecodingGpt2Argument(hparams.n_vocab,
enc.encoder['<|endoftext|>'],
enc.encoder['<|endoftext|>'],
length + 2, decoder_args, top_k,
top_p, temperature)
ckpt_dict = {}
for var in tf.trainable_variables():
ckpt_dict[var.name] = var
decoding_vars = tf.trainable_variables()
op_output = ft_gpt_op(decoding_vars, decoding_args, batch_size,
start_ids, min_start_length, max_start_length,
attention_mask)
generated = 0
while nsamples == 0 or generated < nsamples:
op_out = sess.run(op_output)
for i in range(batch_size):
generated += 1
text = enc.decode(op_out[i])
print("=" * 40 + " SAMPLE " + str(generated) + " " + "=" * 40)
print(text)
from_data = np.random.randn(batch_size, seq_len, encoder_hidden_dim)
from_tensor = tf.convert_to_tensor(from_data, dtype=tf_datatype)
memory_sequence_length = np.random.randint(1,
max_seq_len + 1,
size=batch_size).astype(
np.int32)
embedding_table = np.random.randn(vocab_size, decoder_hidden_dim).astype(
np_datatype) # a [vocab_size, decoder_hidden_dim] table
mask = np.random.randint(2, size=(batch_size, seq_len, seq_len))
attention_mask = tf.convert_to_tensor(mask, dtype=tf_datatype)
encoder_args = TransformerArgument(batch_size=batch_size,
beam_width=1,
head_num=encoder_head_num,
size_per_head=encoder_size_per_head,
num_layer=encoder_num_layer,
max_seq_len=max_seq_len,
dtype=tf_datatype)
decoding_args = DecodingArgument(batch_size=batch_size,
beam_width=beam_width,
head_num=decoder_head_num,
size_per_head=decoder_size_per_head,
num_layer=decoder_num_layer,
max_seq_len=max_seq_len,
vocab_size=vocab_size,
start_id=start_of_sentence_id,
end_id=end_of_sentence_id,
encoder_hidden_dim=encoder_head_num *
encoder_size_per_head,
def encoder_sample(args_dict):
print("\n=============== Argument ===============")
for key in args_dict:
print("{}: {}".format(key, args_dict[key]))
print("========================================")
np.random.seed(1)
tf.set_random_seed(1)
batch_size = args_dict['batch_size']
num_layer = args_dict['num_layer']
max_seq_len = args_dict['max_seq_len']
avg_seq_len = args_dict['avg_seq_len']
head_num = args_dict['head_number']
size_per_head = args_dict['size_per_head']
tf_datatype = tf.float32
np_datatype = np.float32
atol_threshold = 3e-5
int8_mode = args_dict['int8_mode']
allow_gemm_test = True if args_dict['allow_gemm_test'].lower() == "true" else False
if args_dict['data_type'] == "fp16":
tf_datatype = tf.float16
np_datatype = np.float16
atol_threshold = 3e-2
hidden_dim = head_num * size_per_head
sequence_length = np.random.randint(1, max_seq_len + 1, size=batch_size)
if avg_seq_len != -1:
# This means we use "remove_padding" and set other average sequence length
sequence_length = np.ones(batch_size) * avg_seq_len
else:
sequence_length = np.ones(batch_size) * (max_seq_len / 2)
sequence_length = sequence_length.astype(np.int32)
from_data = np.random.randn(batch_size, max_seq_len, hidden_dim)
from_tensor = tf.convert_to_tensor(from_data, dtype=tf_datatype)
attention_mask = build_sequence_mask(sequence_length, num_heads=head_num, maximum_length=max_seq_len, dtype=tf_datatype)
encoder_args = TransformerArgument(beam_width=1,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
remove_padding=False,
int8_mode=int8_mode,
allow_gemm_test=allow_gemm_test)
eff_encoder_args = copy.deepcopy(encoder_args)
eff_encoder_args.remove_padding = True
tf_encoder_result = tf_encoder(input_tensor=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask)
encoder_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
encoder_variables_dict = {}
for v in encoder_vars:
encoder_variables_dict[v.name] = v
op_encoder_result = op_encoder(inputs=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask,
encoder_vars_dict=encoder_variables_dict,
sequence_length=sequence_length)
eff_encoder_result = op_encoder(inputs=from_tensor,
encoder_args=eff_encoder_args,
attention_mask=attention_mask,
encoder_vars_dict=encoder_variables_dict,
sequence_length=sequence_length)
'''
Because FasterTransformer skip some computation for the padding parts,
if we do not mask these parts, the cross check result would be wrong.
'''
tf_encoder_result = tf_encoder_result * tf.expand_dims(tf.sequence_mask(sequence_length, maxlen=max_seq_len, dtype=tf_datatype), axis=-1)
op_encoder_result = op_encoder_result * tf.expand_dims(tf.sequence_mask(sequence_length, maxlen=max_seq_len, dtype=tf_datatype), axis=-1)
eff_encoder_result = eff_encoder_result * tf.expand_dims(tf.sequence_mask(sequence_length, maxlen=max_seq_len, dtype=tf_datatype), axis=-1)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
for idx, name in enumerate(encoder_variables_dict):
print((str(idx) + " " + str(name) + " " +
str(encoder_variables_dict[name].shape)) + " " + str(encoder_variables_dict[name].dtype))
print("#################################")
tf_encoder_result_val = sess.run(tf_encoder_result)
op_encoder_result_val = sess.run(op_encoder_result)
eff_encoder_result_val = sess.run(eff_encoder_result)
cross_check("Encoder TF v.s. FT with tensor input", tf_encoder_result_val, op_encoder_result_val, atol_threshold)
cross_check("Encoder TF v.s. EFF-FT with tensor input", tf_encoder_result_val, eff_encoder_result_val, atol_threshold)
op_diff = abs(tf_encoder_result_val.reshape([-1]) - op_encoder_result_val.reshape([-1]))
eff_diff = abs(tf_encoder_result_val.reshape([-1]) - eff_encoder_result_val.reshape([-1]))
max_diff = max(op_diff.max(), eff_diff.max())
ite = 50
def _cond(from_tensor):
return tf.constant(True)
def _ft_body(from_tensor):
op_encoder_result = op_encoder(inputs=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask,
encoder_vars_dict=encoder_variables_dict,
sequence_length=sequence_length)
return op_encoder_result
def _eff_body(from_tensor):
eff_encoder_result = op_encoder(inputs=from_tensor,
encoder_args=eff_encoder_args,
attention_mask=attention_mask,
encoder_vars_dict=encoder_variables_dict,
sequence_length=sequence_length)
return eff_encoder_result
def _tf_body(from_tensor):
tf_encoder_result = tf_encoder(input_tensor=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask)
return tf_encoder_result
tf_while_tensor = tf.while_loop(_cond,
_tf_body,
loop_vars=[from_tensor],
back_prop=False,
maximum_iterations=ite)
ft_while_tensor = tf.while_loop(_cond,
_ft_body,
loop_vars=[from_tensor],
back_prop=False,
maximum_iterations=ite)
eff_while_tensor = tf.while_loop(_cond,
_eff_body,
loop_vars=[from_tensor],
back_prop=False,
maximum_iterations=ite)
if args_dict['test_time'] == 1:
# tf_time = time_test(sess, tf_encoder_result, ite)
# ft_time = time_test(sess, op_encoder_result, ite)
# eff_time = time_test(sess, eff_encoder_result, ite)
# Using while loop to run 'ite' times to ignore the overheads of memory copy and model preprocess.
# We use these times as the profiling results.
tf_while_time = time_test(sess, tf_while_tensor, 1) / ite # while_loop has run ite times
time.sleep(60)
ft_while_time = time_test(sess, ft_while_tensor, 1) / ite # while_loop has run ite times
time.sleep(60)
eff_while_time = time_test(sess, eff_while_tensor, 1) / ite # while_loop has run ite times
time.sleep(60)
ft_type = args_dict['data_type'].upper()
if int8_mode != 0:
ft_type = "INT8-v{}".format(int8_mode)
# print("[INFO] batch_size {} max_seq_len {} precision {} {} layer TF-time {:6.2f} ms".format(batch_size, max_seq_len, args_dict['data_type'].upper(), num_layer, tf_time))
# print("[INFO] batch_size {} max_seq_len {} precision {} {} layer FT-OP-time {:6.2f} ms".format(batch_size, max_seq_len, ft_type, num_layer, ft_time))
# print("[INFO] batch_size {} max_seq_len {} precision {} {} layer EFF-OP-time {:6.2f} ms".format(batch_size, max_seq_len, ft_type, num_layer, eff_time))
print("[INFO] batch_size {} max_seq_len {} precision {} {} layer TF-while-time {:6.2f} ms ( {} iterations)".format(batch_size, max_seq_len, args_dict['data_type'].upper(), num_layer, tf_while_time, ite))
print("[INFO] batch_size {} max_seq_len {} precision {} {} layer FT-OP-while-time {:6.2f} ms ( {} iterations)".format(batch_size, max_seq_len, ft_type, num_layer, ft_while_time, ite))
print("[INFO] batch_size {} max_seq_len {} precision {} {} layer EFF-OP-while-time {:6.2f} ms ( {} iterations)".format(batch_size, max_seq_len, ft_type, num_layer, eff_while_time, ite))
if args_dict['thread_num'] > 1:
# Multi-threading demonstration
thread_list = []
thread_num = args_dict['thread_num']
def run():
ft_while_time = time_test(sess, ft_while_tensor, 1) / ite # while_loop has run ite times
print("[INFO] batch_size {} max_seq_len {} {} layer FT-OP-while-time {:6.2f} ms with {} threads".format(batch_size,
max_seq_len, num_layer, ft_while_time, thread_num))
for i in range(thread_num):
thread_list.append(threading.Thread(target=run, name="RunFT"))
for t in thread_list:
t.start()
for t in thread_list:
t.join()
return max_diff
size_per_head = args.size_per_head
num_layer = args.num_layer
hidden_dim = head_num * size_per_head
memory_hidden_dim = args.memory_hidden_dim
vocab_size = args.vocab_size
tf_datatype = tf.float32
np_datatype = np.float32
if args.data_type == "fp16":
tf_datatype = tf.float16
np_datatype = np.float16
decoder_args = TransformerArgument(
beam_width=beam_width,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
kernel_init_range=kernel_initializer_range,
bias_init_range=bias_initializer_range,
fuse_qkv=True,
memory_hidden_dim=memory_hidden_dim)
decoding_args = DecodingBeamsearchArgument(vocab_size,
start_of_sentence_id,
end_of_sentence_id, max_seq_len,
decoder_args, 0.0)
embedding_table = np.random.randn(vocab_size, hidden_dim).astype(
np_datatype) * 0.01 # a [vocab_size, hidden_dim] table
embedding_table = tf.convert_to_tensor(embedding_table)
memory, memory_sequence_length = generate_encoder_result(
batch_size, max_seq_len, memory_hidden_dim, tf_datatype)
size=batch_size).astype(np.int32)
if avg_seq_len != -1 and remove_padding == True:
# This means we use "remove_padding" and set a smaller average sequence length
sequence_length = np.ones(batch_size) * avg_seq_len
from_data = np.random.randn(batch_size, max_seq_len, hidden_dim)
from_tensor = tf.convert_to_tensor(from_data, dtype=tf_datatype)
attention_mask = build_sequence_mask(sequence_length,
num_heads=head_num,
maximum_length=max_seq_len,
dtype=tf_datatype)
encoder_args = TransformerArgument(beam_width=1,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
remove_padding=remove_padding)
tf_encoder_result = tf_encoder(input_tensor=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask)
encoder_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
encoder_variables_dict = {}
for v in encoder_vars:
encoder_variables_dict[v.name] = v
op_encoder_result = op_encoder(inputs=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask,
max_seq_len + 1,
size=batch_size).astype(
np.int32)
memory_sequence_length[np.random.randint(0, batch_size)] = max_seq_len
embedding_table = np.random.randn(vocab_size, decoder_hidden_dim).astype(
np_datatype
) * initializer_range # a [vocab_size, decoder_hidden_dim] table
attention_mask = build_sequence_mask(memory_sequence_length,
num_heads=encoder_head_num,
maximum_length=max_seq_len,
dtype=tf_datatype)
encoder_args = TransformerArgument(beam_width=1,
head_num=encoder_head_num,
size_per_head=encoder_size_per_head,
num_layer=encoder_num_layer,
dtype=tf_datatype,
remove_padding=remove_padding)
decoder_args = TransformerArgument(
beam_width=beam_width,
head_num=decoder_head_num,
size_per_head=decoder_size_per_head,
num_layer=decoder_num_layer,
dtype=tf_datatype,
kernel_init_range=kernel_initializer_range,
bias_init_range=bias_initializer_range,
fuse_qkv=False)
decoding_args = DecodingBeamsearchArgument(vocab_size,
start_of_sentence_id,
def encoder_sample(args_dict):
print("\n=============== Argument ===============")
for key in args_dict:
print("{}: {}".format(key, args_dict[key]))
print("========================================")
np.random.seed(1)
tf.set_random_seed(1)
batch_size = args_dict['batch_size']
num_layer = args_dict['num_layer']
max_seq_len = args_dict['max_seq_len']
avg_seq_len = args_dict['avg_seq_len']
head_num = args_dict['head_number']
size_per_head = args_dict['size_per_head']
remove_padding = True if args_dict['remove_padding'].lower(
) == "true" else False
tf_datatype = tf.float32
np_datatype = np.float32
atol_threshold = 3e-5
int8_mode = args_dict['int8_mode']
allow_gemm_test = True if args_dict['allow_gemm_test'].lower(
) == "true" else False
if args_dict['data_type'] == "fp16":
tf_datatype = tf.float16
np_datatype = np.float16
atol_threshold = 3e-2
hidden_dim = head_num * size_per_head
sequence_length = np.random.randint(1, max_seq_len + 1, size=batch_size)
if avg_seq_len != -1 and remove_padding == True:
# This means we use "remove_padding" and set a smaller average sequence length
sequence_length = np.ones(batch_size) * avg_seq_len
else:
sequence_length = np.ones(batch_size) * (max_seq_len / 2)
sequence_length = sequence_length.astype(np.int32)
from_data = np.random.randn(batch_size, max_seq_len, hidden_dim)
from_tensor = tf.convert_to_tensor(from_data, dtype=tf_datatype)
attention_mask = build_sequence_mask(sequence_length,
num_heads=head_num,
maximum_length=max_seq_len,
dtype=tf_datatype)
encoder_args = TransformerArgument(beam_width=1,
head_num=head_num,
size_per_head=size_per_head,
num_layer=num_layer,
dtype=tf_datatype,
remove_padding=remove_padding,
int8_mode=int8_mode,
allow_gemm_test=allow_gemm_test)
tf_encoder_result = tf_encoder(input_tensor=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask)
encoder_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
encoder_variables_dict = {}
for v in encoder_vars:
encoder_variables_dict[v.name] = v
op_encoder_result = op_encoder(inputs=from_tensor,
encoder_args=encoder_args,
attention_mask=attention_mask,
encoder_vars_dict=encoder_variables_dict,
sequence_length=sequence_length)
'''
Because FasterTransformer skip some computation for the padding parts,
if we do not mask these parts, the cross check result would be wrong.
'''
tf_encoder_result = tf_encoder_result * tf.expand_dims(tf.sequence_mask(
sequence_length, maxlen=max_seq_len, dtype=tf_datatype),
axis=-1)
op_encoder_result = op_encoder_result * tf.expand_dims(tf.sequence_mask(
sequence_length, maxlen=max_seq_len, dtype=tf_datatype),
axis=-1)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
for idx, name in enumerate(encoder_variables_dict):
print((str(idx) + " " + str(name) + " " +
str(encoder_variables_dict[name].shape)) + " " +
str(encoder_variables_dict[name].dtype))
print("#################################")
tf_encoder_result_val = sess.run(tf_encoder_result)
op_encoder_result_val = sess.run(op_encoder_result)
cross_check("Encoder TF v.s. FT with tensor input",
tf_encoder_result_val, op_encoder_result_val,
atol_threshold)
'''
Use the numpy array as inputs of FasterTransformer OP.
This method require more time for the op initialization (especially for FP16),
but the inference time would be little faster than using tensor as input.
'''
encoder_variables_dict_2 = {}
for var, val in zip(encoder_vars, sess.run(encoder_vars)):
encoder_variables_dict_2[var.name] = val
# op_encoder_result_2 = op_encoder(inputs=from_tensor,
# encoder_args=encoder_args,
# attention_mask=attention_mask,
# encoder_vars_dict=encoder_variables_dict_2,
# sequence_length=sequence_length)
# op_encoder_result_val_2 = sess.run(op_encoder_result_2)
# cross_check("Encoder TF v.s. FT with numpy input", tf_encoder_result_val,
# op_encoder_result_val_2, atol_threshold)
if args_dict['test_time'] == 1:
ite = 50
tf_time = time_test(sess, tf_encoder_result, ite)
op_time = time_test(sess, op_encoder_result, ite)
# op_time_2 = time_test(sess, op_encoder_result_2, ite)
print(
"[INFO] batch_size {} max_seq_len {} {} layer TF-time {:6.2f} ms"
.format(batch_size, max_seq_len, num_layer, tf_time))
print(
"[INFO] batch_size {} max_seq_len {} {} layer FT-OP-tensor-time {:6.2f} ms"
.format(batch_size, max_seq_len, num_layer, op_time))
# print("[INFO] batch_size {} max_seq_len {} {} layer FT-OP-numpy-time {:6.2f} ms".format(batch_size, max_seq_len, num_layer, op_time_2))
return (tf_encoder_result_val.reshape([-1]) -
op_encoder_result_val.reshape([-1])).max()
def sample_model(model_name='124M',
nsamples=1,
batch_size=1,
length=12,
temperature=1,
top_k=4,
top_p=0,
models_dir='models',
data_type='fp32'):
"""Run the sample_model.
:model_name=124M : String, which model to use
:nsamples=0 : Number of samples to return, if 0, continues to
generate samples indefinately.
:batch_size=1 : Number of batches (only affects speed/memory).
:length=None : Number of tokens in generated text, if None (default), is
determined by model hyperparameters
:temperature=1 : Float value controlling randomness in boltzmann
distribution. Lower temperature results in less random completions. As the
temperature approaches zero, the model will become deterministic and
repetitive. Higher temperature results in more random completions.
:top_k=4 : Integer value controlling diversity. 1 means only 1 word is
considered for each step (token), resulting in deterministic completions,
while 40 means 40 words are considered at each step. 0 (default) is a
special setting meaning no restrictions. 40 generally is a good value.
:models_dir : path to parent folder containing model subfolders
(i.e. contains the <model_name> folder)
"""
models_dir = os.path.expanduser(os.path.expandvars(models_dir))
enc = encoder.get_encoder(model_name, models_dir)
hparams = HParams(n_vocab=0, n_ctx=1024, n_embd=768, n_head=12, n_layer=12)
with open(os.path.join(models_dir, model_name, 'hparams.json')) as f:
hparams.override_from_dict(json.load(f))
if length is None:
length = hparams.n_ctx
elif length > hparams.n_ctx:
raise ValueError("Can't get samples longer than window size: %s" %
hparams.n_ctx)
# start_ids has shape [batch_size, start_len].flatten()
# start_ids = [15496, 11, 616, 3290, 468,
# 15496, 11, 616, 3290, 469,
# 15496, 11, 616, 3290, 470,
# 15496, 11, 616, 3290, 471]
start_ids = [enc.encoder['<|endoftext|>'] for i in range(batch_size)]
with tf.Session(graph=tf.Graph()) as sess:
saver = tf.train.import_meta_graph("{}/{}/model.ckpt.meta".format(
models_dir, model_name))
print("[INFO] restore the model {}/{}".format(models_dir, model_name))
saver.restore(sess,
("{}/{}/model.ckpt".format(models_dir, model_name)))
if data_type == 'fp32':
tf_data_type = tf.float32
elif data_type == 'fp16':
tf_data_type = tf.float16
else:
assert (False)
decoder_args = TransformerArgument(beam_width=1,
head_num=hparams.n_head,
size_per_head=hparams.n_embd //
hparams.n_head,
num_layer=hparams.n_layer,
dtype=tf_data_type,
kernel_init_range=0.00,
bias_init_range=0.00)
decoding_args = DecodingGpt2Argument(hparams.n_vocab,
enc.encoder['<|endoftext|>'],
enc.encoder['<|endoftext|>'],
length + 2, decoder_args, top_k,
top_p, temperature)
ckpt_dict = {}
for var in tf.trainable_variables():
ckpt_dict[var.name] = var
decoding_vars = tf.trainable_variables()
op_output = ft_gpt2_op(decoding_vars, decoding_args, batch_size,
start_ids)
generated = 0
while nsamples == 0 or generated < nsamples:
print("[INFO] FT op time: {}".format(
time_test(sess, op_output, iterations=5, warmup=True)))
op_out = sess.run(op_output)
for i in range(batch_size):
generated += 1
text = enc.decode(op_out[i][1:])
print("=" * 40 + " SAMPLE " + str(generated) + " " + "=" * 40)
print(text)
def translate_sample(args_dict):
print("\n=============== Argument ===============")
for key in args_dict:
print("{}: {}".format(key, args_dict[key]))
print("========================================")
np.random.seed(1)
tf.set_random_seed(1)
random.seed(1)
start_of_sentence_id = 1
end_of_sentence_id = 2
kernel_initializer_range = 0.02
bias_initializer_range = 0.02
batch_size = args_dict['batch_size']
beam_width = args_dict['beam_width']
max_seq_len = args_dict['max_seq_len']
encoder_head_num = args_dict['encoder_head_number']
encoder_size_per_head = args_dict['encoder_size_per_head']
decoder_head_num = args_dict['decoder_head_number']
decoder_size_per_head = args_dict['decoder_size_per_head']
encoder_num_layer = args_dict['encoder_num_layer']
decoder_num_layer = args_dict['decoder_num_layer']
encoder_hidden_dim = encoder_head_num * encoder_size_per_head
decoder_hidden_dim = decoder_head_num * decoder_size_per_head
tf_datatype = tf.float32
if args_dict['data_type'] == "fp16":
tf_datatype = tf.float16
beam_search_diversity_rate = args_dict['beam_search_diversity_rate']
sampling_topk = args_dict['sampling_topk']
sampling_topp = args_dict['sampling_topp']
source_inputter = WordEmbedder("source_vocabulary",
embedding_size=encoder_hidden_dim,
dtype=tf_datatype)
target_inputter = WordEmbedder("target_vocabulary",
embedding_size=decoder_hidden_dim,
dtype=tf_datatype)
inputter = ExampleInputter(source_inputter, target_inputter)
inputter.initialize({
"source_vocabulary": args_dict['source_vocabulary'],
"target_vocabulary": args_dict['target_vocabulary']
})
vocab_size = target_inputter.vocabulary_size
source_file = args_dict['source']
is_remove_padding = True if args_dict['remove_padding'].lower(
) == "true" else False
encoder_args = TransformerArgument(
beam_width=1,
head_num=encoder_head_num,
size_per_head=encoder_size_per_head,
num_layer=encoder_num_layer,
dtype=tf_datatype,
kernel_init_range=kernel_initializer_range,
bias_init_range=bias_initializer_range,
remove_padding=is_remove_padding)
decoder_args = TransformerArgument(
beam_width=beam_width,
head_num=decoder_head_num,
size_per_head=decoder_size_per_head,
num_layer=decoder_num_layer,
dtype=tf_datatype,
kernel_init_range=kernel_initializer_range,
bias_init_range=bias_initializer_range,
memory_hidden_dim=encoder_head_num * encoder_size_per_head)
decoder_args_2 = copy.deepcopy(decoder_args) # for beam search
decoder_args_2.__dict__ = copy.deepcopy(decoder_args.__dict__)
decoder_args_2.beam_width = 1 # for sampling
decoding_beamsearch_args = DecodingBeamsearchArgument(
vocab_size, start_of_sentence_id, end_of_sentence_id, max_seq_len,
decoder_args, beam_search_diversity_rate)
decoding_sampling_args = DecodingSamplingArgument(
vocab_size, start_of_sentence_id, end_of_sentence_id, max_seq_len,
decoder_args_2, sampling_topk, sampling_topp)
with tf.variable_scope("transformer/encoder", reuse=tf.AUTO_REUSE):
dataset = inputter.make_inference_dataset(source_file, batch_size)
iterator = dataset.make_initializable_iterator()
source = iterator.get_next()
source_embedding = source_inputter.make_inputs(source)
source_embedding = tf.cast(source_embedding, tf_datatype)
memory_sequence_length = source["length"]
tf_encoder_result = tf_encoder_opennmt(
source_embedding,
encoder_args,
sequence_length=memory_sequence_length)
encoder_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
encoder_variables_dict = {}
for v in encoder_vars:
encoder_variables_dict[v.name] = v
ft_encoder_result = ft_encoder_opennmt(
inputs=source_embedding,
encoder_args=encoder_args,
encoder_vars_dict=encoder_variables_dict,
sequence_length=memory_sequence_length)
tf_encoder_result = tf.reshape(tf_encoder_result,
tf.shape(source_embedding))
ft_encoder_result = tf.reshape(ft_encoder_result,
tf.shape(source_embedding))
with tf.variable_scope("transformer/decoder", reuse=tf.AUTO_REUSE):
target_inputter.build()
target_vocab_rev = target_inputter.vocabulary_lookup_reverse()
### TF BeamSearch Decoding ###
tf_beamsearch_target_ids, tf_beamsearch_target_length, _, _, _ = tf_beamsearch_decoding(
tf_encoder_result,
memory_sequence_length,
target_inputter.embedding,
decoding_beamsearch_args,
decoder_type=0)
# tf_beamsearch_target_tokens: [batch_size, beam_width, seq_len]
tf_beamsearch_target_tokens = target_vocab_rev.lookup(
tf.cast(tf_beamsearch_target_ids, tf.int64))
tf_beamsearch_target_length = tf.minimum(
tf_beamsearch_target_length + 1,
tf.shape(tf_beamsearch_target_ids)[-1])
### end of TF BeamSearch Decoding ###
### TF Sampling Decoding ###
tf_sampling_target_ids, tf_sampling_target_length = tf_sampling_decoding(
tf_encoder_result,
memory_sequence_length,
target_inputter.embedding,
decoding_sampling_args,
decoder_type=0)
# tf_sampling_target_tokens: [batch_size, seq_len]
tf_sampling_target_tokens = target_vocab_rev.lookup(
tf.cast(tf_sampling_target_ids, tf.int64))
tf_sampling_target_length = tf.minimum(
tf_sampling_target_length + 1,
tf.shape(tf_sampling_target_ids)[-1])
### end of TF BeamSearch Decoding ###
### OP BeamSearch Decoder ###
op_decoder_beamsearch_target_ids, op_decoder_beamsearch_target_length, _, _, _ = tf_beamsearch_decoding(
tf_encoder_result,
memory_sequence_length,
target_inputter.embedding,
decoding_beamsearch_args,
decoder_type=1)
# op_decoder_beamsearch_target_tokens: [batch_size, beam_width, seq_len]
op_decoder_beamsearch_target_tokens = target_vocab_rev.lookup(
tf.cast(op_decoder_beamsearch_target_ids, tf.int64))
op_decoder_beamsearch_target_length = tf.minimum(
op_decoder_beamsearch_target_length + 1,
tf.shape(op_decoder_beamsearch_target_ids)[-1])
### end of OP BeamSearch Decoder ###
### OP Sampling Decoder ###
op_decoder_sampling_target_ids, op_decoder_sampling_target_length = tf_sampling_decoding(
tf_encoder_result,
memory_sequence_length,
target_inputter.embedding,
decoding_sampling_args,
decoder_type=1)
op_decoder_sampling_target_tokens = target_vocab_rev.lookup(
tf.cast(op_decoder_sampling_target_ids, tf.int64))
op_decoder_sampling_target_length = tf.minimum(
op_decoder_sampling_target_length + 1,
tf.shape(op_decoder_sampling_target_ids)[-1])
### end of OP BeamSearch Decoder ###
### Prepare Decoding variables for FasterTransformer ###
all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
decoder_var_start_id = 0
while all_vars[decoder_var_start_id].name.find(
"transformer/decoder") == -1:
decoder_var_start_id += 1
decoder_variables = all_vars[
decoder_var_start_id +
1:] # decoder_var_start_id + 1 means skip the embedding table
### OP BeamSearch Decoding ###
op_beamsearch_target_ids, op_beamsearch_target_length, _, _, _ = op_beamsearch_decoding(
ft_encoder_result, memory_sequence_length, target_inputter.embedding,
decoder_variables, decoding_beamsearch_args)
op_beamsearch_target_tokens = target_vocab_rev.lookup(
tf.cast(op_beamsearch_target_ids, tf.int64))
op_beamsearch_target_length = tf.minimum(
op_beamsearch_target_length + 1,
tf.shape(op_beamsearch_target_ids)[-1])
### end of OP BeamSearch Decoding ###
### OP Sampling Decoding ###
op_sampling_target_ids, op_sampling_target_length = op_sampling_decoding(
ft_encoder_result, memory_sequence_length, target_inputter.embedding,
decoder_variables, decoding_sampling_args)
op_sampling_target_tokens = target_vocab_rev.lookup(
tf.cast(op_sampling_target_ids, tf.int64))
op_sampling_target_length = tf.minimum(
op_sampling_target_length + 1,
tf.shape(op_sampling_target_ids)[-1])
### end of OP Sampling Decoding ###
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
time_args = args_dict['test_time']
class TranslationResult(object):
def __init__(self, token_op, length_op, name):
self.token_op = token_op
self.length_op = length_op
self.name = name
self.file_name = name + ".txt"
self.token_list = []
self.length_list = []
self.batch_num = 0
self.execution_time = 0.0 # seconds
self.sentence_num = 0
self.bleu_score = None
translation_result_list = []
if time_args.find("0") != -1:
translation_result_list.append(
TranslationResult(tf_beamsearch_target_tokens,
tf_beamsearch_target_length,
"tf-decoding-beamsearch"))
if time_args.find("1") != -1:
translation_result_list.append(
TranslationResult(op_decoder_beamsearch_target_tokens,
op_decoder_beamsearch_target_length,
"op-decoder-beamsearch"))
if time_args.find("2") != -1:
translation_result_list.append(
TranslationResult(op_beamsearch_target_tokens,
op_beamsearch_target_length,
"op-decoding-beamsearch"))
if time_args.find("3") != -1:
translation_result_list.append(
TranslationResult(tf_sampling_target_tokens,
tf_sampling_target_length,
"tf-decoding-sampling"))
if time_args.find("4") != -1:
translation_result_list.append(
TranslationResult(op_decoder_sampling_target_tokens,
op_decoder_sampling_target_length,
"op-decoder-sampling"))
if time_args.find("5") != -1:
translation_result_list.append(
TranslationResult(op_sampling_target_tokens,
op_sampling_target_length,
"op-decoding-sampling"))
float_var_list = []
half_var_list = []
for var in tf.global_variables()[:-1]:
if var.dtype.base_dtype == tf.float32:
float_var_list.append(var)
elif var.dtype.base_dtype == tf.float16:
half_var_list.append(var)
if (len(translation_result_list) == 0):
print("[WARNING] No put any test cases.")
cuda_profiler = cudaProfiler()
cuda_profiler.start()
for i in range(len(translation_result_list)):
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
sess.run(tf.tables_initializer())
sess.run(iterator.initializer)
if (len(float_var_list) > 0):
float_saver = tf.train.Saver(float_var_list)
float_saver.restore(sess, "translation/ckpt/model.ckpt-500000")
if (len(half_var_list) > 0):
half_saver = tf.train.Saver(half_var_list)
half_saver.restore(sess,
"translation/ckpt/fp16_model.ckpt-500000")
t1 = datetime.now()
while True:
try:
batch_tokens, batch_length = sess.run([
translation_result_list[i].token_op,
translation_result_list[i].length_op
])
for tokens, length in zip(batch_tokens, batch_length):
if translation_result_list[i].name.find(
"beamsearch") != -1:
translation_result_list[i].token_list.append(
b" ".join(tokens[0][:length[0] -
2]).decode("UTF-8"))
else:
translation_result_list[i].token_list.append(
b" ".join(tokens[:length - 2]).decode("UTF-8"))
translation_result_list[i].batch_num += 1
except tf.errors.OutOfRangeError:
break
t2 = datetime.now()
time_sum = (t2 - t1).total_seconds()
translation_result_list[i].execution_time = time_sum
with open(translation_result_list[i].file_name, "w") as file_b:
for s in translation_result_list[i].token_list:
file_b.write(s)
file_b.write("\n")
ref_file_path = "./.ref_file.txt"
os.system("head -n %d %s > %s" %
(len(translation_result_list[i].token_list),
args_dict['target'], ref_file_path))
translation_result_list[i].bleu_score = bleu_score(
translation_result_list[i].file_name, ref_file_path)
os.system("rm {}".format(ref_file_path))
time.sleep(60)
cuda_profiler.stop()
for t in translation_result_list:
print(
"[INFO] {} translates {} batches taking {:.2f} sec to translate {} tokens, BLEU score: {:.2f}, {:.0f} tokens/sec."
.format(t.name, t.batch_num, t.execution_time,
t.bleu_score.sys_len, t.bleu_score.score,
t.bleu_score.sys_len / t.execution_time))
return translation_result_list
