adding automatic mixed precision training support

2019-10-14 12:34:10 +00:00 · 2019-10-14 12:34:10 +00:00 · 909fa60601
commit 909fa60601
parent b888058e4e
1 changed files with 4 additions and 0 deletions
--- a/DeepSpeech.py
+++ b/DeepSpeech.py
@ -446,6 +446,10 @@ def train():

    # Building the graph
    optimizer = create_optimizer()
+    
+    # Enable mixed precision training
+    if FLAGS.automatic_mixed_precision:
+        optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(optimizer)
    gradients, loss, non_finite_files = get_tower_results(iterator, optimizer, dropout_rates)

    # Average tower gradients across GPUs