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