def chpt_to_dict_arrays():
"""
Convert a checkpoint into a dictionary of numpy arrays
for later use in TensorRT NMT sample.
git clone https://github.com/tensorflow/nmt.git
"""
sys.path.append("./nmt")
from nmt.nmt import add_arguments, create_hparams
from nmt import attention_model
from nmt import model_helper
from nmt.nmt import create_or_load_hparams
from nmt import utils
from nmt import model as nmt_model
nmt_parser = argparse.ArgumentParser()
add_arguments(nmt_parser)
FLAGS, unparsed = nmt_parser.parse_known_args()
default_hparams = create_hparams(FLAGS)
hparams = create_or_load_hparams(FLAGS.out_dir,
default_hparams,
FLAGS.hparams_path,
save_hparams=False)
print(hparams)
model_creator = None
if not hparams.attention:
model_creator = nmt_model.Model
elif hparams.attention_architecture == "standard":
model_creator = attention_model.AttentionModel
else:
raise ValueError("Unknown model architecture")
infer_model = model_helper.create_infer_model(model_creator,
hparams,
scope=None)
params = {}
print("\nFound the following trainable variables:")
with tf.Session(graph=infer_model.graph,
config=utils.misc_utils.get_config_proto()) as sess:
loaded_infer_model = model_helper.load_model(infer_model.model,
FLAGS.ckpt, sess, "infer")
variables = tf.trainable_variables()
for v in variables:
params[v.name] = v.eval(session=sess)
print("{0} {1}".format(v.name, params[v.name].shape))
params["forget_bias"] = hparams.forget_bias
return params
def testTrain(self):
"""Test the training loop is functional with basic hparams."""
nmt_parser = argparse.ArgumentParser()
nmt.add_arguments(nmt_parser)
FLAGS, unparsed = nmt_parser.parse_known_args()
_update_flags(FLAGS, "nmt_train_test")
default_hparams = nmt.create_hparams(FLAGS)
train_fn = train.train
nmt.run_main(FLAGS, default_hparams, train_fn, None)
def do_start_inference(out_dir, hparams):
# Silence all outputs
#os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
global current_stdout
current_stdout = sys.stdout
sys.stdout = open(os.devnull, "w")
# Modified autorun from nmt.py (bottom of the file)
# We want to use original argument parser (for validation, etc)
nmt_parser = argparse.ArgumentParser()
nmt.add_arguments(nmt_parser)
# But we have to hack settings from our config in there instead of commandline options
flags, unparsed = nmt_parser.parse_known_args(
['--' + k + '=' + str(v) for k, v in hparams.items()])
# And now we can run TF with modified arguments
#tf.app.run(main=nmt.main, argv=[os.getcwd() + '\nmt\nmt\nmt.py'] + unparsed)
# Add output (model) folder to flags
flags.out_dir = out_dir
# Make hparams
hparams = nmt.create_hparams(flags)
## Train / Decode
if not tf.gfile.Exists(flags.out_dir):
nmt.utils.print_out("# Model folder (out_dir) doesn't exist")
sys.exit()
# Load hparams from model folder
hparams = nmt.create_or_load_hparams(flags.out_dir,
hparams,
flags.hparams_path,
save_hparams=True)
# Choose checkpoint (provided with hparams or last one)
if not flags.ckpt:
flags.ckpt = tf.train.latest_checkpoint(flags.out_dir)
# Create model
if not hparams.attention:
model_creator = nmt.inference.nmt_model.Model
elif hparams.attention_architecture == "standard":
model_creator = nmt.inference.attention_model.AttentionModel
elif hparams.attention_architecture in ["gnmt", "gnmt_v2"]:
model_creator = nmt.inference.gnmt_model.GNMTModel
else:
raise ValueError("Unknown model architecture")
infer_model = nmt.inference.model_helper.create_infer_model(
model_creator, hparams, None)
return (infer_model, flags, hparams)
def testInference(self):
"""Test inference is function with basic hparams."""
nmt_parser = argparse.ArgumentParser()
nmt.add_arguments(nmt_parser)
FLAGS, unparsed = nmt_parser.parse_known_args()
_update_flags(FLAGS, "nmt_train_infer")
# Train one step so we have a checkpoint.
FLAGS.num_train_steps = 1
default_hparams = nmt.create_hparams(FLAGS)
train_fn = train.train
nmt.run_main(FLAGS, default_hparams, train_fn, None)
# Update FLAGS for inference.
FLAGS.inference_input_file = ("nmt/testdata/" "iwslt15.tst2013.100.en")
FLAGS.inference_output_file = os.path.join(FLAGS.out_dir, "output")
FLAGS.inference_ref_file = ("nmt/testdata/" "iwslt15.tst2013.100.vi")
default_hparams = nmt.create_hparams(FLAGS)
inference_fn = inference.inference
nmt.run_main(FLAGS, default_hparams, None, inference_fn)
def init(FLAGS):
global hparams, ckpt, infer_model
FLAGS.num_units = 128
FLAGS.share_vocab = True
default_hparams = nmt.create_hparams(FLAGS)
# Load hparams.
hparams = nmt.create_or_load_hparams(FLAGS.out_dir,
default_hparams,
FLAGS.hparams_path,
save_hparams=False)
ckpt = tf.train.latest_checkpoint(FLAGS.out_dir)
infer_model = inference.inference_m(ckpt, hparams, scope=None)
def setup(self, flags):
# Model output directory
out_dir = flags.out_dir
if out_dir and not tf.gfile.Exists(out_dir):
tf.gfile.MakeDirs(out_dir)
# Load hparams.
default_hparams = create_hparams(flags)
loaded_hparams = False
if flags.ckpt: # Try to load hparams from the same directory as ckpt
ckpt_dir = os.path.dirname(flags.ckpt)
ckpt_hparams_file = os.path.join(ckpt_dir, "hparams")
if tf.gfile.Exists(ckpt_hparams_file) or flags.hparams_path:
# Note: for some reason this will create an empty "best_bleu" directory and copy vocab files
hparams = create_or_load_hparams(ckpt_dir,
default_hparams,
flags.hparams_path,
save_hparams=False)
loaded_hparams = True
assert loaded_hparams
# GPU device
config_proto = utils.get_config_proto(
allow_soft_placement=True,
num_intra_threads=hparams.num_intra_threads,
num_inter_threads=hparams.num_inter_threads)
utils.print_out("# Devices visible to TensorFlow: %s" %
repr(tf.Session(config=config_proto).list_devices()))
# Inference indices (inference_indices is broken, but without setting it to None we'll crash)
hparams.inference_indices = None
# Create the graph
model_creator = get_model_creator(hparams)
infer_model = model_helper.create_infer_model(model_creator,
hparams,
scope=None)
sess, loaded_infer_model = start_sess_and_load_model(
infer_model, flags.ckpt, hparams)
# Parameters needed by TF GNMT
self.hparams = hparams
self.infer_model = infer_model
self.sess = sess
self.loaded_infer_model = loaded_infer_model
def setup_inference_parameters(out_dir, hparams):
# Print output on stdout while temporarily sending other output to /dev/null
global current_stdout
current_stdout = sys.stdout
sys.stdout = open(os.devnull, "w")
nmt_parser = ap.ArgumentParser()
nmt.add_arguments(nmt_parser)
# Get settingds from configuration file
flags, unparsed = nmt_parser.parse_known_args(
['--' + key + '=' + str(value) for key, value in hparams.items()])
# Add output (model) folder to flags
flags.out_dir = out_dir
## Exit if model folder doesn't exist
if not tf.gfile.Exists(flags.out_dir):
nmt.utils.print_out("# Model folder (out_dir) doesn't exist")
sys.exit()
# Load hyper parameters (hparams) from model folder
hparams = nmt.create_hparams(flags)
hparams = nmt.create_or_load_hparams(flags.out_dir,
hparams,
flags.hparams_path,
save_hparams=True)
# Choose checkpoint (provided with hparams or last one)
if not flags.ckpt:
flags.ckpt = tf.train.latest_checkpoint(flags.out_dir)
# Create model
if not hparams.attention:
model_creator = nmt.inference.nmt_model.Model
elif hparams.attention_architecture == "standard":
model_creator = nmt.inference.attention_model.AttentionModel
elif hparams.attention_architecture in ["gnmt", "gnmt_v2"]:
model_creator = nmt.inference.gnmt_model.GNMTModel
else:
raise ValueError("Unknown model architecture")
infer_model = nmt.inference.model_helper.create_infer_model(
model_creator, hparams, None)
return (infer_model, flags, hparams)
def do_start_inference(out_dir, hparams):
global current_stdout
current_stdout = sys.stdout
sys.stdout = open(os.devnull, "w")
# Modified autorun from nmt.py (bottom of the file)
# We want to use original argument parser
nmt_parser = argparse.ArgumentParser()
nmt.add_arguments(nmt_parser)
flags, unparsed = nmt_parser.parse_known_args(['--'+k+'='+str(v) for k,v in hparams.items()])
# Add output (model) folder to flags
flags.out_dir = out_dir
# Make hparams
hparams = nmt.create_hparams(flags)
if not tf.gfile.Exists(flags.out_dir):
nmt.utils.print_out("# Model folder (out_dir) doesn't exist")
sys.exit()
# Load hparams from model folder
hparams = nmt.create_or_load_hparams(flags.out_dir, hparams, flags.hparams_path, save_hparams=True)
if not flags.ckpt:
flags.ckpt = tf.train.latest_checkpoint(flags.out_dir)
if not hparams.attention:
model_creator = nmt.inference.nmt_model.Model
elif hparams.attention_architecture == "standard":
model_creator = nmt.inference.attention_model.AttentionModel
elif hparams.attention_architecture in ["gnmt", "gnmt_v2"]:
model_creator = nmt.inference.gnmt_model.GNMTModel
else:
raise ValueError("Unknown model architecture")
infer_model = nmt.inference.model_helper.create_infer_model(model_creator, hparams, None)
return (infer_model, flags, hparams)
def generate_reply(input_text, flags):
# Format data
tokenized_text = tokenize_text(input_text)
infer_data = [tokenized_text]
# Load hparams.
jobid = flags.jobid
default_hparams = nmt.create_hparams(flags)
hparams = nmt.create_or_load_hparams(model_dir,
default_hparams,
flags.hparams_path,
save_hparams=(jobid == 0))
# Load checkpoint
ckpt = tf.train.latest_checkpoint(model_dir)
# Inference
model_creator = attention_model.AttentionModel
# Create model
scope = None
infer_model = model_helper.create_infer_model(model_creator, hparams,
scope)
with tf.Session(graph=infer_model.graph,
config=misc_utils.get_config_proto()) as sess:
model = model_helper.load_model(infer_model.model, ckpt, sess, "infer")
sess.run(infer_model.iterator.initializer,
feed_dict={
infer_model.src_placeholder: infer_data,
infer_model.batch_size_placeholder:
hparams.infer_batch_size
})
# Decode
nmt_outputs, infer_summary = model.decode(sess)
# get text translation(reply as a chatbot)
assert nmt_outputs.shape[0] == 1
translation = nmt_utils.get_translation(
nmt_outputs,
sent_id=0,
tgt_eos=hparams.eos,
subword_option=hparams.subword_option)
return translation.decode("utf-8")
def predicted(file_path, output_path, model_path, src, tgt):
add_arg = nmt.add_arguments
nmt_parser = argparse.ArgumentParser()
add_arg(nmt_parser)
FLAGS, unparsed = nmt_parser.parse_known_args([
'--src', src, '--tgt', tgt, '--vocab_prefix',
'nmt/testdata/voc_50w', '--out_dir', model_path,
'--inference_input_file', file_path, '--inference_output_file',
output_path
])
default_hparams = nmt.create_hparams(FLAGS)
train_fn = nmt.train.train
inference_fn = nmt.inference.inference
nmt.run_main(FLAGS, default_hparams, train_fn, inference_fn)
with open(output_path, 'r', encoding="utf-8") as file:
prediction = file.read().replace('\n', '')
if is_contain_chinese(prediction) == True:
result = str(prediction).replace(' ', '')
else:
result = str(prediction)
return result
def setup(self, flags):
# Model output directory
out_dir = flags.out_dir
if out_dir and not tf.gfile.Exists(out_dir):
tf.gfile.MakeDirs(out_dir)
# Load hparams.
default_hparams = create_hparams(flags)
loaded_hparams = False
if flags.ckpt: # Try to load hparams from the same directory as ckpt
ckpt_dir = os.path.dirname(flags.ckpt)
ckpt_hparams_file = os.path.join(ckpt_dir, "hparams")
if tf.gfile.Exists(ckpt_hparams_file) or flags.hparams_path:
# Note: for some reason this will create an empty "best_bleu" directory and copy vocab files
hparams = create_or_load_hparams(ckpt_dir, default_hparams, flags.hparams_path, save_hparams=False)
loaded_hparams = True
assert loaded_hparams
# GPU device
config_proto = utils.get_config_proto(
allow_soft_placement=True,
num_intra_threads=hparams.num_intra_threads,
num_inter_threads=hparams.num_inter_threads)
utils.print_out(
"# Devices visible to TensorFlow: %s"
% repr(tf.Session(config=config_proto).list_devices()))
# Inference indices (inference_indices is broken, but without setting it to None we'll crash)
hparams.inference_indices = None
# Create the graph
model_creator = get_model_creator(hparams)
infer_model = model_helper.create_infer_model(model_creator, hparams, scope=None)
sess, loaded_infer_model = start_sess_and_load_model(infer_model, flags.ckpt,
hparams)
# FIXME (bryce): Set to False to disable inference from frozen graph and run fast again
if True:
frozen_graph = None
with infer_model.graph.as_default():
output_node_names = ['hash_table_Lookup_1/LookupTableFindV2']
other_node_names = ['MakeIterator', 'IteratorToStringHandle', 'init_all_tables', 'NoOp', 'dynamic_seq2seq/decoder/NoOp']
frozen_graph = tf.graph_util.convert_variables_to_constants(sess,
tf.get_default_graph().as_graph_def(),
output_node_names=output_node_names + other_node_names)
# FIXME (bryce): Uncomment this block to enable tensorRT convert
from tensorflow.python.compiler.tensorrt import trt_convert as trt
converter = trt.TrtGraphConverter(input_graph_def=frozen_graph, nodes_blacklist=(output_node_names),
is_dynamic_op=True, max_batch_size=hparams.infer_batch_size,
max_beam_size=hparams.beam_width, max_src_seq_len=hparams.src_max_len)
frozen_graph = converter.convert()
with tf.Graph().as_default():
tf.graph_util.import_graph_def(frozen_graph, name="")
sess = tf.Session(graph=tf.get_default_graph(),
config=utils.get_config_proto(
num_intra_threads=hparams.num_intra_threads,
num_inter_threads=hparams.num_inter_threads)
)
iterator = iterator_utils.BatchedInput(
initializer=tf.get_default_graph().get_operation_by_name(infer_model.iterator.initializer.name),
source=tf.get_default_graph().get_tensor_by_name(infer_model.iterator.source.name),
target_input=None,
target_output=None,
source_sequence_length=tf.get_default_graph().get_tensor_by_name(infer_model.iterator.source_sequence_length.name),
target_sequence_length=None)
infer_model = model_helper.InferModel(
graph=tf.get_default_graph(),
model=infer_model.model,
src_placeholder=tf.get_default_graph().get_tensor_by_name(infer_model.src_placeholder.name),
batch_size_placeholder=tf.get_default_graph().get_tensor_by_name(infer_model.batch_size_placeholder.name),
iterator=iterator)
# Parameters needed by TF GNMT
self.hparams = hparams
self.infer_model = infer_model
self.sess = sess
self.loaded_infer_model = loaded_infer_model