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