def compile_train_fn(model, learning_rate=2e-4): """ Build the CTC training routine for speech models. Args: model: A keras model (built=True) instance Returns: train_fn (theano.function): Function that takes in acoustic inputs, and updates the model. Returns network outputs and ctc cost """ logger.info("Building train_fn") acoustic_input = model.inputs[0] network_output = model.outputs[0] output_lens = K.placeholder(ndim=1, dtype='int32') label = K.placeholder(ndim=1, dtype='int32') label_lens = K.placeholder(ndim=1, dtype='int32') network_output = network_output.dimshuffle((1, 0, 2)) ctc_cost = ctc.cpu_ctc_th(network_output, output_lens, label, label_lens).mean() trainable_vars = model.trainable_weights optimizer = SGD(nesterov=True, lr=learning_rate, momentum=0.9, clipnorm=100) updates = optimizer.get_updates(trainable_vars, [], ctc_cost) train_fn = K.function( [acoustic_input, output_lens, label, label_lens, K.learning_phase()], [network_output, ctc_cost], updates=updates) return train_fn
def make_train_fn(self, model, learning_rate, momentum, optimizer, regmask_start_val, regmask_anneal_episodes): assert (regmask_start_val is None) == (regmask_anneal_episodes is None) if optimizer == 'sgd': opt = SGD(lr=learning_rate, momentum=momentum) elif optimizer == 'adam': opt = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-8) else: raise ValueError('Unrecognized optimizer: {}'.format(optimizer)) state_input = model.input action_input = Input(batch_shape=(None, ), dtype=np.int32, name='action_input') reward_input = Input(batch_shape=(None, ), dtype=np.float32, name='reward_input') episode = tf.placeholder(dtype=tf.int32, shape=(), name='episode') action_probas = model(state_input) indexer = tf.stack( [tf.range(0, tf.shape(action_input)[0]), action_input], axis=1) selected_action_probas = tf.gather_nd(params=action_probas, indices=indexer) mean_entropy = K.mean( -K.sum(action_probas * K.log(action_probas + 1e-6), axis=1)) objectives = reward_input * K.log(selected_action_probas) entropy_loss_decay_coefficient = K.exp( -self.entropy_boost_decay * tf.cast(episode, dtype=tf.float32)) undiscounted_entropy_loss = -K.log(mean_entropy + self.EPSILON) entropy_loss = entropy_loss_decay_coefficient * self.entropy_boost * undiscounted_entropy_loss regmask_loss = _get_regmask_loss(regmask_start_val, regmask_anneal_episodes, self.state_shape, self.action_space.n, model.layers[0].kernel, episode) loss = (-K.sum(objectives) + entropy_loss + regmask_loss) updates = opt.get_updates(loss=loss, params=model.trainable_weights) # --- Statistics for monitoring purposes only grads_overall = opt.get_gradients(loss=loss, params=model.trainable_weights) grads_entropy = opt.get_gradients(loss=entropy_loss, params=model.trainable_weights) # We assume a linear model with weight-matrix and bias-vector assert len(grads_overall) == len(grads_entropy) == 2 W_grad_length_overall = K.sqrt(K.sum(K.square(grads_overall[0]))) W_grad_length_entropy = K.sqrt(K.sum(K.square(grads_entropy[0]))) # weight_variance = K.var(model.trainable_weights[0]) # ------------------------------- train_fn = K.function( inputs=[state_input, action_input, reward_input, episode], outputs=[ entropy_loss, undiscounted_entropy_loss, W_grad_length_overall, W_grad_length_entropy, regmask_loss ], updates=updates) return train_fn