def model_fn(features, labels, mode, params): conformer_model = conformer.Model( kernel_regularizer = None, bias_regularizer = None, **config ) decoder_config = malaya_speech.config.conformer_base_decoder_config transducer_model = transducer.rnn.Model( conformer_model, vocabulary_size = subwords.vocab_size, **decoder_config ) v = tf.expand_dims(features['inputs'], -1) z = tf.zeros((tf.shape(features['targets'])[0], 1), dtype = tf.int64) c = tf.concat([z, features['targets']], axis = 1) logits = transducer_model([v, c], training = True) cost = transducer.loss.rnnt_loss( logits = logits, labels = features['targets'], label_length = features['targets_length'][:, 0], logit_length = features['inputs_length'][:, 0] // conformer_model.conv_subsampling.time_reduction_factor, ) mean_error = tf.reduce_mean(cost) loss = mean_error tf.identity(loss, 'train_loss') variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) init_checkpoint = 'asr-base-conformer-transducer/model.ckpt-150000' assignment_map, initialized_variable_names = train.get_assignment_map_from_checkpoint( variables, init_checkpoint ) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) if mode == tf.estimator.ModeKeys.TRAIN: train_op = train.optimizer.optimize_loss( loss, tf.train.AdamOptimizer, parameters['optimizer_params'], learning_rate_scheduler, summaries = ['learning_rate', 'loss_scale'], larc_params = parameters.get('larc_params', None), loss_scaling = parameters.get('loss_scaling', 1.0), loss_scaling_params = parameters.get('loss_scaling_params', None), ) estimator_spec = tf.estimator.EstimatorSpec( mode = mode, loss = loss, train_op = train_op ) elif mode == tf.estimator.ModeKeys.EVAL: estimator_spec = tf.estimator.EstimatorSpec( mode = tf.estimator.ModeKeys.EVAL, loss = loss ) return estimator_spec
def model_fn(features, labels, mode, params): learning_rate = 1e-5 init_checkpoint = '../deep-speaker/out/vggvox.ckpt' Y = tf.cast(features['targets'][:, 0], tf.int32) model = deep_speaker.model.Model( features['inputs'], num_class = 10, mode = 'train' ) logits = model.logits loss = tf.reduce_mean( tf.nn.sparse_softmax_cross_entropy_with_logits( logits = logits, labels = Y ) ) tf.identity(loss, 'train_loss') accuracy = tf.metrics.accuracy( labels = Y, predictions = tf.argmax(logits, axis = 1) ) tf.identity(accuracy[1], name = 'train_accuracy') variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) variables = [v for v in variables if 'prediction' not in v.name] assignment_map, initialized_variable_names = train.get_assignment_map_from_checkpoint( variables, init_checkpoint ) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) if mode == tf.estimator.ModeKeys.TRAIN: global_step = tf.train.get_or_create_global_step() optimizer = tf.train.AdamOptimizer(learning_rate) train_op = optimizer.minimize(loss, global_step = global_step) estimator_spec = tf.estimator.EstimatorSpec( mode = mode, loss = loss, train_op = train_op ) elif mode == tf.estimator.ModeKeys.EVAL: estimator_spec = tf.estimator.EstimatorSpec( mode = tf.estimator.ModeKeys.EVAL, loss = loss, eval_metric_ops = {'accuracy': accuracy}, ) return estimator_spec
def model_fn(features, labels, mode, params): model = Model(features['combined'], [features['y'], features['noise']]) loss = model.cost tf.identity(loss, 'total_loss') tf.summary.scalar('total_loss', loss) for i in range(len(model.loss)): tf.identity(model.loss[i], f'loss_{i}') tf.summary.scalar(f'loss_{i}', model.loss[i]) global_step = tf.train.get_or_create_global_step() variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) assignment_map, initialized_variable_names = train.get_assignment_map_from_checkpoint( variables, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32) learning_rate = tf.train.polynomial_decay( learning_rate, global_step, epochs, end_learning_rate=1e-7, power=1.0, cycle=False, ) tf.summary.scalar('learning_rate', learning_rate) if mode == tf.estimator.ModeKeys.TRAIN: optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate) train_op = optimizer.minimize(loss, global_step=global_step) estimator_spec = tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op) elif mode == tf.estimator.ModeKeys.EVAL: estimator_spec = tf.estimator.EstimatorSpec( mode=tf.estimator.ModeKeys.EVAL, loss=loss) return estimator_spec