def session_config(params):
optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L1,
do_function_inlining=True)
graph_options = tf.GraphOptions(optimizer_options=optimizer_options)
config = tf.ConfigProto(allow_soft_placement=True,
graph_options=graph_options)
if distribute.is_distributed_training_mode():
config.gpu_options.visible_device_list = str(distribute.local_rank())
elif params.device_list:
device_str = ",".join([str(i) for i in params.device_list])
config.gpu_options.visible_device_list = device_str
return config
def session_config(params):
optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L1,
do_function_inlining=True)
graph_options = tf.GraphOptions(optimizer_options=optimizer_options)
config = tf.ConfigProto(allow_soft_placement=True,
graph_options=graph_options)
from tensorflow.python.client import device_lib
n_gpus = sum(1 for d in device_lib.list_local_devices() if d.device_type == 'GPU')
if n_gpus > 1:
params.device_list = list(range(n_gpus))
if distribute.is_distributed_training_mode():
config.gpu_options.visible_device_list = str(distribute.local_rank())
elif params.device_list:
device_str = ",".join([str(i) for i in params.device_list])
config.gpu_options.visible_device_list = device_str
return config
def get_training_input(filenames, params):
""" Get input for training stage
:param filenames: A list contains [source_filenames, target_filenames]
:param params: Hyper-parameters
:returns: A dictionary of pair <Key, Tensor>
"""
with tf.device("/cpu:0"):
src_dataset = tf.data.TextLineDataset(filenames[0])
tgt_dataset = tf.data.TextLineDataset(filenames[1])
dataset = tf.data.Dataset.zip((src_dataset, tgt_dataset))
if distribute.is_distributed_training_mode():
dataset = dataset.shard(distribute.size(), distribute.rank())
dataset = dataset.shuffle(params.buffer_size)
dataset = dataset.repeat()
# Split string
dataset = dataset.map(
lambda src, tgt:
(tf.string_split([src]).values, tf.string_split([tgt]).values),
num_parallel_calls=params.num_threads)
# Append <eos> symbol
dataset = dataset.map(
lambda src, tgt:
(tf.concat([src, [tf.constant(params.eos)]], axis=0),
tf.concat([tgt, [tf.constant(params.eos)]], axis=0)),
num_parallel_calls=params.num_threads)
# Convert to dictionary
dataset = dataset.map(lambda src, tgt: {
"source": src,
"target": tgt,
"source_length": tf.shape(src),
"target_length": tf.shape(tgt)
},
num_parallel_calls=params.num_threads)
# Create iterator
iterator = dataset.make_one_shot_iterator()
features = iterator.get_next()
# Create lookup table
src_table = tf.contrib.lookup.index_table_from_tensor(
tf.constant(params.vocabulary["source"]),
default_value=params.mapping["source"][params.unk])
tgt_table = tf.contrib.lookup.index_table_from_tensor(
tf.constant(params.vocabulary["target"]),
default_value=params.mapping["target"][params.unk])
# String to index lookup
features["source"] = src_table.lookup(features["source"])
features["target"] = tgt_table.lookup(features["target"])
# Batching
features = batch_examples(features,
params.batch_size,
params.max_length,
params.mantissa_bits,
shard_multiplier=len(params.device_list),
length_multiplier=params.length_multiplier,
constant=params.constant_batch_size,
num_threads=params.num_threads)
# Convert to int32
features["source"] = tf.to_int32(features["source"])
features["target"] = tf.to_int32(features["target"])
features["source_length"] = tf.to_int32(features["source_length"])
features["target_length"] = tf.to_int32(features["target_length"])
features["source_length"] = tf.squeeze(features["source_length"], 1)
features["target_length"] = tf.squeeze(features["target_length"], 1)
return features
def get_training_input(filenames, params):
""" Get input for training stage
:param filenames: A list contains [source_filenames, target_filenames]
:param params: Hyper-parameters
:returns: A dictionary of pair <Key, Tensor>
"""
with tf.device("/cpu:0"):
datasets = []
for filename in filenames:
datasets.append(tf.data.TextLineDataset(filename))
dataset = tf.data.Dataset.zip(tuple(datasets))
if distribute.is_distributed_training_mode():
dataset = dataset.shard(distribute.size(), distribute.rank())
dataset = dataset.shuffle(params.buffer_size)
dataset = dataset.repeat()
# Split string
dataset = dataset.map(
lambda *x: [tf.string_split([y]).values for y in x],
num_parallel_calls=params.num_threads)
# Append <eos> symbol
dataset = dataset.map(
lambda *x:
[tf.concat([y, [tf.constant(params.eos)]], axis=0) for y in x],
num_parallel_calls=params.num_threads)
def convert_to_dict(src, tgt, *x):
res = {}
res["source"] = src
res["source_length"] = tf.shape(src)
res["target"] = tgt
res["target_length"] = tf.shape(tgt)
for i, v in enumerate(x):
res["mt_%d" % i] = v
res["mt_length_%d" % i] = tf.shape(v)
return res
# Convert to dictionary
dataset = dataset.map(convert_to_dict,
num_parallel_calls=params.num_threads)
# Create iterator
iterator = dataset.make_one_shot_iterator()
features = iterator.get_next()
# Create lookup table
src_table = tf.contrib.lookup.index_table_from_tensor(
tf.constant(params.vocabulary["source"]),
default_value=params.mapping["source"][params.unk])
tgt_table = tf.contrib.lookup.index_table_from_tensor(
tf.constant(params.vocabulary["target"]),
default_value=params.mapping["target"][params.unk])
# String to index lookup
features["source"] = src_table.lookup(features["source"])
features["target"] = tgt_table.lookup(features["target"])
for i in range(len(filenames) - 2):
features["mt_%d" % i] = tgt_table.lookup(features["mt_%d" % i])
# Batching
features = batch_examples(features,
params.batch_size,
params.max_length,
params.mantissa_bits,
shard_multiplier=len(params.device_list),
length_multiplier=params.length_multiplier,
constant=params.constant_batch_size,
num_threads=params.num_threads)
# Convert to int32
features["source"] = tf.to_int32(features["source"])
features["target"] = tf.to_int32(features["target"])
features["source_length"] = tf.to_int32(features["source_length"])
features["target_length"] = tf.to_int32(features["target_length"])
features["source_length"] = tf.squeeze(features["source_length"], 1)
features["target_length"] = tf.squeeze(features["target_length"], 1)
for i in range(len(filenames) - 2):
features["mt_%d" % i] = tf.to_int32(features["mt_%d" % i])
features["mt_length_%d" % i] = tf.to_int32(
features["mt_length_%d" % i])
features["mt_length_%d" % i] = tf.squeeze(
features["mt_length_%d" % i], 1)
return features
