From d55cbe6dfd7adaf42f677f1263fe4f2124d0779c Mon Sep 17 00:00:00 2001 From: josh Date: Tue, 2 Apr 2019 01:02:50 +0200 Subject: [PATCH] accidentally merged out the transfer learning! recovering it now --- DeepSpeech.py | 291 ++++++++++++++++++++++++++++++++------------------ 1 file changed, 185 insertions(+), 106 deletions(-) diff --git a/DeepSpeech.py b/DeepSpeech.py index 0a2656bf..85f63e63 100755 --- a/DeepSpeech.py +++ b/DeepSpeech.py @@ -12,11 +12,13 @@ import evaluate import numpy as np import progressbar import shutil +import tempfile import tensorflow as tf import traceback from ds_ctcdecoder import ctc_beam_search_decoder, Scorer from six.moves import zip, range +from tensorflow.contrib.lite.python import tflite_convert from tensorflow.python.tools import freeze_graph from util.audio import audiofile_to_input_vector from util.config import Config, initialize_globals @@ -27,12 +29,6 @@ from util.logging import log_info, log_error, log_debug, log_warn from util.preprocess import preprocess from util.text import Alphabet -#TODO: remove once fully switched to 1.13 -try: - from tensorflow.contrib.lite.python import tflite_convert # 1.12 -except ImportError: - from tensorflow.lite.python import tflite_convert # 1.13 - # Graph Creation # ============== @@ -88,84 +84,90 @@ def BiRNN(batch_x, seq_length, dropout, reuse=False, batch_size=None, n_steps=-1 # clipped RELU activation and dropout. # 1st layer - b1 = variable_on_worker_level('b1', [Config.n_hidden_1], tf.zeros_initializer()) - h1 = variable_on_worker_level('h1', [Config.n_input + 2*Config.n_input*Config.n_context, Config.n_hidden_1], tf.contrib.layers.xavier_initializer()) - layer_1 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(batch_x, h1), b1)), FLAGS.relu_clip) - layer_1 = tf.nn.dropout(layer_1, rate=dropout[0]) - layers['layer_1'] = layer_1 + with tf.name_scope("LAYER_1"): + b1 = variable_on_worker_level('b1', [Config.n_hidden_1], tf.zeros_initializer()) + h1 = variable_on_worker_level('h1', [Config.n_input + 2*Config.n_input*Config.n_context, Config.n_hidden_1], tf.contrib.layers.xavier_initializer()) + layer_1 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(batch_x, h1), b1)), FLAGS.relu_clip) + layer_1 = tf.nn.dropout(layer_1, (1.0 - dropout[0])) + layers['layer_1'] = layer_1 # 2nd layer - b2 = variable_on_worker_level('b2', [Config.n_hidden_2], tf.zeros_initializer()) - h2 = variable_on_worker_level('h2', [Config.n_hidden_1, Config.n_hidden_2], tf.contrib.layers.xavier_initializer()) - layer_2 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_1, h2), b2)), FLAGS.relu_clip) - layer_2 = tf.nn.dropout(layer_2, rate=dropout[1]) - layers['layer_2'] = layer_2 + with tf.name_scope("LAYER_2"): + b2 = variable_on_worker_level('b2', [Config.n_hidden_2], tf.zeros_initializer()) + h2 = variable_on_worker_level('h2', [Config.n_hidden_1, Config.n_hidden_2], tf.contrib.layers.xavier_initializer()) + layer_2 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_1, h2), b2)), FLAGS.relu_clip) + layer_2 = tf.nn.dropout(layer_2, (1.0 - dropout[1])) + layers['layer_2'] = layer_2 # 3rd layer - b3 = variable_on_worker_level('b3', [Config.n_hidden_3], tf.zeros_initializer()) - h3 = variable_on_worker_level('h3', [Config.n_hidden_2, Config.n_hidden_3], tf.contrib.layers.xavier_initializer()) - layer_3 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_2, h3), b3)), FLAGS.relu_clip) - layer_3 = tf.nn.dropout(layer_3, rate=dropout[2]) - layers['layer_3'] = layer_3 + with tf.name_scope("LAYER_3"): + b3 = variable_on_worker_level('b3', [Config.n_hidden_3], tf.zeros_initializer()) + h3 = variable_on_worker_level('h3', [Config.n_hidden_2, Config.n_hidden_3], tf.contrib.layers.xavier_initializer()) + layer_3 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_2, h3), b3)), FLAGS.relu_clip) + layer_3 = tf.nn.dropout(layer_3, (1.0 - dropout[2])) + layers['layer_3'] = layer_3 # Now we create the forward and backward LSTM units. # Both of which have inputs of length `n_cell_dim` and bias `1.0` for the forget gate of the LSTM. - # Forward direction cell: - if not tflite: - fw_cell = tf.contrib.rnn.LSTMBlockFusedCell(Config.n_cell_dim, reuse=reuse) - layers['fw_cell'] = fw_cell - else: - fw_cell = tf.nn.rnn_cell.LSTMCell(Config.n_cell_dim, reuse=reuse) + with tf.name_scope("LAYER_LSTM"): + # Forward direction cell: + if not tflite: + fw_cell = tf.contrib.rnn.LSTMBlockFusedCell(Config.n_cell_dim, reuse=reuse) + layers['fw_cell'] = fw_cell + else: + fw_cell = tf.nn.rnn_cell.LSTMCell(Config.n_cell_dim, reuse=reuse) - # `layer_3` is now reshaped into `[n_steps, batch_size, 2*n_cell_dim]`, - # as the LSTM RNN expects its input to be of shape `[max_time, batch_size, input_size]`. - layer_3 = tf.reshape(layer_3, [n_steps, batch_size, Config.n_hidden_3]) - if tflite: - # Generated StridedSlice, not supported by NNAPI - #n_layer_3 = [] - #for l in range(layer_3.shape[0]): - # n_layer_3.append(layer_3[l]) - #layer_3 = n_layer_3 + # `layer_3` is now reshaped into `[n_steps, batch_size, 2*n_cell_dim]`, + # as the LSTM RNN expects its input to be of shape `[max_time, batch_size, input_size]`. + layer_3 = tf.reshape(layer_3, [n_steps, batch_size, Config.n_hidden_3]) + if tflite: + # Generated StridedSlice, not supported by NNAPI + #n_layer_3 = [] + #for l in range(layer_3.shape[0]): + # n_layer_3.append(layer_3[l]) + #layer_3 = n_layer_3 + + # Unstack/Unpack is not supported by NNAPI + layer_3 = tf.unstack(layer_3, n_steps) - # Unstack/Unpack is not supported by NNAPI - layer_3 = tf.unstack(layer_3, n_steps) + # We parametrize the RNN implementation as the training and inference graph + # need to do different things here. + if not tflite: + output, output_state = fw_cell(inputs=layer_3, dtype=tf.float32, sequence_length=seq_length, initial_state=previous_state) + else: + output, output_state = tf.nn.static_rnn(fw_cell, layer_3, previous_state, tf.float32) + output = tf.concat(output, 0) - # We parametrize the RNN implementation as the training and inference graph - # need to do different things here. - if not tflite: - output, output_state = fw_cell(inputs=layer_3, dtype=tf.float32, sequence_length=seq_length, initial_state=previous_state) - else: - output, output_state = tf.nn.static_rnn(fw_cell, layer_3, previous_state, tf.float32) - output = tf.concat(output, 0) - - # Reshape output from a tensor of shape [n_steps, batch_size, n_cell_dim] - # to a tensor of shape [n_steps*batch_size, n_cell_dim] - output = tf.reshape(output, [-1, Config.n_cell_dim]) - layers['rnn_output'] = output - layers['rnn_output_state'] = output_state + # Reshape output from a tensor of shape [n_steps, batch_size, n_cell_dim] + # to a tensor of shape [n_steps*batch_size, n_cell_dim] + output = tf.reshape(output, [-1, Config.n_cell_dim]) + layers['rnn_output'] = output + layers['rnn_output_state'] = output_state # Now we feed `output` to the fifth hidden layer with clipped RELU activation and dropout - b5 = variable_on_worker_level('b5', [Config.n_hidden_5], tf.zeros_initializer()) - h5 = variable_on_worker_level('h5', [Config.n_cell_dim, Config.n_hidden_5], tf.contrib.layers.xavier_initializer()) - layer_5 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(output, h5), b5)), FLAGS.relu_clip) - layer_5 = tf.nn.dropout(layer_5, rate=dropout[5]) - layers['layer_5'] = layer_5 + with tf.name_scope("LAYER_5"): + b5 = variable_on_worker_level('b5', [Config.n_hidden_5], tf.zeros_initializer()) + h5 = variable_on_worker_level('h5', [Config.n_cell_dim, Config.n_hidden_5], tf.contrib.layers.xavier_initializer()) + layer_5 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(output, h5), b5)), FLAGS.relu_clip) + layer_5 = tf.nn.dropout(layer_5, (1.0 - dropout[5])) + layers['layer_5'] = layer_5 # Now we apply the weight matrix `h6` and bias `b6` to the output of `layer_5` # creating `n_classes` dimensional vectors, the logits. - b6 = variable_on_worker_level('b6', [Config.n_hidden_6], tf.zeros_initializer()) - h6 = variable_on_worker_level('h6', [Config.n_hidden_5, Config.n_hidden_6], tf.contrib.layers.xavier_initializer()) - layer_6 = tf.add(tf.matmul(layer_5, h6), b6) - layers['layer_6'] = layer_6 + with tf.name_scope("LAYER_6"): + b6 = variable_on_worker_level('b6', [Config.n_hidden_6], tf.zeros_initializer()) + h6 = variable_on_worker_level('h6', [Config.n_hidden_5, Config.n_hidden_6], tf.contrib.layers.xavier_initializer()) + layer_6 = tf.add(tf.matmul(layer_5, h6), b6) + layers['layer_6'] = layer_6 - # Finally we reshape layer_6 from a tensor of shape [n_steps*batch_size, n_hidden_6] - # to the slightly more useful shape [n_steps, batch_size, n_hidden_6]. - # Note, that this differs from the input in that it is time-major. - layer_6 = tf.reshape(layer_6, [n_steps, batch_size, Config.n_hidden_6], name="raw_logits") - layers['raw_logits'] = layer_6 + # Finally we reshape layer_6 from a tensor of shape [n_steps*batch_size, n_hidden_6] + # to the slightly more useful shape [n_steps, batch_size, n_hidden_6]. + # Note, that this differs from the input in that it is time-major. + layer_6 = tf.reshape(layer_6, [n_steps, batch_size, Config.n_hidden_6], name="raw_logits") + layers['raw_logits'] = layer_6 - # Output shape: [n_steps, batch_size, n_hidden_6] + # Output shape: [n_steps, batch_size, n_hidden_6] return layer_6, layers @@ -192,7 +194,10 @@ def calculate_mean_edit_distance_and_loss(model_feeder, tower, dropout, reuse): logits, _ = BiRNN(batch_x, batch_seq_len, dropout, reuse) # Compute the CTC loss using TensorFlow's `ctc_loss` - total_loss = tf.nn.ctc_loss(labels=batch_y, inputs=logits, sequence_length=batch_seq_len) + total_loss = tf.nn.ctc_loss(labels=batch_y, + inputs=logits, + sequence_length=batch_seq_len, + ignore_longer_outputs_than_inputs=True) # Calculate the average loss across the batch avg_loss = tf.reduce_mean(total_loss) @@ -234,7 +239,7 @@ def create_optimizer(): # on which all operations within the tower execute. # For example, all operations of 'tower 0' could execute on the first GPU `tf.device('/gpu:0')`. -def get_tower_results(model_feeder, optimizer, dropout_rates): +def get_tower_results(model_feeder, optimizer, dropout_rates, drop_source_layers): r''' With this preliminary step out of the way, we can for each GPU introduce a tower for which's batch we calculate and return the optimization gradients @@ -267,8 +272,25 @@ def get_tower_results(model_feeder, optimizer, dropout_rates): # Retain tower's avg losses tower_avg_losses.append(avg_loss) - # Compute gradients for model parameters using tower's mini-batch - gradients = optimizer.compute_gradients(avg_loss) + # # Compute gradients for model parameters using tower's mini-batch + if drop_source_layers == 0: + # train from scratch and update all layers + gradients = optimizer.compute_gradients(avg_loss) + elif FLAGS.fine_tune: + # train from source model and fine-tine + # aka - update all layers + gradients = optimizer.compute_gradients(avg_loss) + else: + # train from source model and freeze old layers + # aka - only update new layers + gradients = optimizer.compute_gradients( + avg_loss, + var_list = [ v for v in tf.trainable_variables() + if any( + layer in v.op.name + for layer in drop_source_layers + )] + ) # Retain tower's gradients tower_gradients.append(gradients) @@ -373,6 +395,22 @@ def train(server=None): If no server provided, it performs single process training. ''' + # The transfer learning approach here need us to supply the layers which we + # want to exclude from the source model. + # Say we want to exclude all layers except for the first one, we can use this: + # + # drop_source_layers=['2', '3', 'lstm', '5', '6'] + # + # If we want to use all layers from the source model except the last one, we use this: + # + # drop_source_layers=['6'] + # + + if FLAGS.drop_source_layers == 0: + drop_source_layers = 0 + else: + drop_source_layers = ['2', '3', 'lstm', '5', '6'][-int(FLAGS.drop_source_layers):] + # Initializing and starting the training coordinator coord = TrainingCoordinator(Config.is_chief) coord.start() @@ -427,7 +465,7 @@ def train(server=None): total_num_replicas=FLAGS.replicas) # Get the data_set specific graph end-points - gradients, loss = get_tower_results(model_feeder, optimizer, dropout_rates) + gradients, loss = get_tower_results(model_feeder, optimizer, dropout_rates, drop_source_layers) # Average tower gradients across GPUs avg_tower_gradients = average_gradients(gradients) @@ -535,6 +573,36 @@ def train(server=None): # Initialize update_progressbar()'s child fields to safe values update_progressbar.pbar = None + ### TRANSFER LEARNING ### + def init_fn(scaffold, session): + if FLAGS.source_model_checkpoint_dir: + print('Initializing from', FLAGS.source_model_checkpoint_dir) + ckpt = tf.train.load_checkpoint(FLAGS.source_model_checkpoint_dir) + variables = list(ckpt.get_variable_to_shape_map().keys()) + for v in tf.global_variables(): + # TRAIN FROM SCRATCH + if drop_source_layers == 0: + print('Loading', v.op.name) + v.load(ckpt.get_tensor(v.op.name), session=session) + else: + # TRAIN FROM SOURCE + if not any(layer in v.op.name for layer in drop_source_layers): + print('Loading', v.op.name) + v.load(ckpt.get_tensor(v.op.name), session=session) + + if drop_source_layers == 0: + init_op=tf.variables_initializer(tf.global_variables()) + else: + init_op=tf.variables_initializer( + [ v for v in tf.global_variables() + if any(layer in v.op.name + for layer in drop_source_layers) + ]) + + scaffold = tf.train.Scaffold(init_op, init_fn=init_fn) + ### TRANSFER LEARNING ### + + # The MonitoredTrainingSession takes care of session initialization, # restoring from a checkpoint, saving to a checkpoint, and closing when done # or an error occurs. @@ -542,12 +610,17 @@ def train(server=None): with tf.train.MonitoredTrainingSession(master='' if server is None else server.target, is_chief=Config.is_chief, hooks=hooks, + scaffold=scaffold, # transfer-learning checkpoint_dir=FLAGS.checkpoint_dir, save_checkpoint_secs=None, # already taken care of by a hook log_step_count_steps=0, # disable logging of steps/s to avoid TF warning in validation sets config=Config.session_config) as session: tf.get_default_graph().finalize() + # add the graph to the "each" summary + step_summary_writers.get('train').add_graph(session.graph) + step_summary_writers.get('dev').add_graph(session.graph) + try: if Config.is_chief: # Retrieving global_step from the (potentially restored) model @@ -663,23 +736,18 @@ def test(): def create_inference_graph(batch_size=1, n_steps=16, tflite=False): - batch_size = batch_size if batch_size > 0 else None # Input tensor will be of shape [batch_size, n_steps, 2*n_context+1, n_input] input_tensor = tf.placeholder(tf.float32, [batch_size, n_steps if n_steps > 0 else None, 2*Config.n_context+1, Config.n_input], name='input_node') seq_length = tf.placeholder(tf.int32, [batch_size], name='input_lengths') - if batch_size <= 0: - # no state management since n_step is expected to be dynamic too (see below) - previous_state = previous_state_c = previous_state_h = None + if not tflite: + previous_state_c = variable_on_worker_level('previous_state_c', [batch_size, Config.n_cell_dim], initializer=None) + previous_state_h = variable_on_worker_level('previous_state_h', [batch_size, Config.n_cell_dim], initializer=None) else: - if not tflite: - previous_state_c = variable_on_worker_level('previous_state_c', [batch_size, Config.n_cell_dim], initializer=None) - previous_state_h = variable_on_worker_level('previous_state_h', [batch_size, Config.n_cell_dim], initializer=None) - else: - previous_state_c = tf.placeholder(tf.float32, [batch_size, Config.n_cell_dim], name='previous_state_c') - previous_state_h = tf.placeholder(tf.float32, [batch_size, Config.n_cell_dim], name='previous_state_h') + previous_state_c = tf.placeholder(tf.float32, [batch_size, Config.n_cell_dim], name='previous_state_c') + previous_state_h = tf.placeholder(tf.float32, [batch_size, Config.n_cell_dim], name='previous_state_h') - previous_state = tf.contrib.rnn.LSTMStateTuple(previous_state_c, previous_state_h) + previous_state = tf.contrib.rnn.LSTMStateTuple(previous_state_c, previous_state_h) no_dropout = [0.0] * 6 @@ -700,23 +768,9 @@ def create_inference_graph(batch_size=1, n_steps=16, tflite=False): # Apply softmax for CTC decoder logits = tf.nn.softmax(logits) - if batch_size <= 0: - if tflite: - raise NotImplementedError('dynamic batch_size does not support tflite nor streaming') - if n_steps > 0: - raise NotImplementedError('dynamic batch_size expect n_steps to be dynamic too') - return ( - { - 'input': input_tensor, - 'input_lengths': seq_length, - }, - { - 'outputs': tf.identity(logits, name='logits'), - }, - layers - ) - new_state_c, new_state_h = layers['rnn_output_state'] + + # Initial zero state if not tflite: zero_state = tf.zeros([batch_size, Config.n_cell_dim], tf.float32) initialize_c = tf.assign(previous_state_c, zero_state) @@ -767,7 +821,7 @@ def export(): tf.reset_default_graph() session = tf.Session(config=Config.session_config) - inputs, outputs, _ = create_inference_graph(batch_size=FLAGS.export_batch_size, n_steps=FLAGS.n_steps, tflite=FLAGS.export_tflite) + inputs, outputs, _ = create_inference_graph(batch_size=1, n_steps=FLAGS.n_steps, tflite=FLAGS.export_tflite) input_names = ",".join(tensor.op.name for tensor in inputs.values()) output_names_tensors = [ tensor.op.name for tensor in outputs.values() if isinstance(tensor, Tensor) ] output_names_ops = [ tensor.name for tensor in outputs.values() if isinstance(tensor, Operation) ] @@ -803,7 +857,7 @@ def export(): os.makedirs(FLAGS.export_dir) def do_graph_freeze(output_file=None, output_node_names=None, variables_blacklist=None): - return freeze_graph.freeze_graph_with_def_protos( + freeze_graph.freeze_graph_with_def_protos( input_graph_def=session.graph_def, input_saver_def=saver.as_saver_def(), input_checkpoint=checkpoint_path, @@ -818,15 +872,39 @@ def export(): if not FLAGS.export_tflite: do_graph_freeze(output_file=output_graph_path, output_node_names=output_names, variables_blacklist='previous_state_c,previous_state_h') else: - frozen_graph = do_graph_freeze(output_node_names=output_names, variables_blacklist='') + temp_fd, temp_freeze = tempfile.mkstemp(dir=FLAGS.export_dir) + os.close(temp_fd) + do_graph_freeze(output_file=temp_freeze, output_node_names=output_names, variables_blacklist='') output_tflite_path = os.path.join(FLAGS.export_dir, output_filename.replace('.pb', '.tflite')) - converter = tf.lite.TFLiteConverter(frozen_graph, input_tensors=inputs.values(), output_tensors=outputs.values()) - converter.post_training_quantize = True - tflite_model = converter.convert() + class TFLiteFlags(): + def __init__(self): + self.graph_def_file = temp_freeze + self.inference_type = 'FLOAT' + self.input_arrays = input_names + self.input_shapes = input_shapes + self.output_arrays = output_names + self.output_file = output_tflite_path + self.output_format = 'TFLITE' - with open(output_tflite_path, 'wb') as fout: - fout.write(tflite_model) + default_empty = [ + 'inference_input_type', + 'mean_values', + 'default_ranges_min', 'default_ranges_max', + 'drop_control_dependency', + 'reorder_across_fake_quant', + 'change_concat_input_ranges', + 'allow_custom_ops', + 'converter_mode', + 'post_training_quantize', + 'dump_graphviz_dir', + 'dump_graphviz_video' + ] + for e in default_empty: + self.__dict__[e] = None + flags = TFLiteFlags() + tflite_convert._convert_model(flags) + os.unlink(temp_freeze) log_info('Exported model for TF Lite engine as {}'.format(os.path.basename(output_tflite_path))) log_info('Models exported at %s' % (FLAGS.export_dir)) @@ -851,6 +929,7 @@ def do_single_file_inference(input_file_path): checkpoint_path = checkpoint.model_checkpoint_path saver.restore(session, checkpoint_path) + session.run(outputs['initialize_state']) features = audiofile_to_input_vector(input_file_path, Config.n_input, Config.n_context)