Make the default values for experimental and non experimental apis match.

PiperOrigin-RevId: 185952648

1 files changed, 22 insertions, 6 deletions

diff --git a/tensorflow/contrib/quantize/python/quantize_graph.py b/tensorflow/contrib/quantize/python/quantize_graph.py
index 0dfe78fd02..5a3a74cec4 100644
--- a/tensorflow/contrib/quantize/python/quantize_graph.py
+++ b/tensorflow/contrib/quantize/python/quantize_graph.py
@@ -69,7 +69,7 @@ def _create_graph(input_graph=None,
         activation_bits=activation_bits)
 
 
-def create_training_graph(input_graph=None, quant_delay=250000):
+def create_training_graph(input_graph=None, quant_delay=0):
   """Rewrites a training input_graph in place for simulated quantization.
 
   The graph has fake quantization ops inserted to simulate the error
@@ -77,6 +77,14 @@ def create_training_graph(input_graph=None, quant_delay=250000):
   the expected behavior of previously held references to nodes and tensors may
   change.
 
+  The default value of quant_delay is suitable for finetuning an already trained
+  floating point model (recommended).
+  If one wants to train a quantized model from scratch, quant_delay should be
+  set to the number of steps it take the floating point model to converge.
+  Quantization will be activated at this point and effectively finetune the
+  model. If quant_delay is not provided when training from scratch, training can
+  often fail.
+
   Args:
     input_graph: The tf.Graph to be transformed.
     quant_delay: Number of steps after which weights and activations are
@@ -93,12 +101,12 @@ def create_training_graph(input_graph=None, quant_delay=250000):
     # Corresponds to case of restoring from a floating point checkpoint
     # In this case, we can freeze the moving mean and variance early on and
     # switch to using them during training. Therefore, freeze_bn_delay is set to
-    # 200000
-    freeze_bn_delay = 200000
+    # 2e5.
+    freeze_bn_delay = int(2e5)
   else:
     # If training from scratch, set freeze_bn_delay to 100 epochs after quant
     # delay. With a batch size of 64, this corresponds to 20000*100=2M steps.
-    freeze_bn_delay = quant_delay + 2000000
+    freeze_bn_delay = quant_delay + int(2e6)
 
   _create_graph(
       input_graph=input_graph,
@@ -129,8 +137,8 @@ def create_eval_graph(input_graph=None):
 def experimental_create_training_graph(input_graph=None,
                                        weight_bits=8,
                                        activation_bits=8,
-                                       quant_delay=250000,
-                                       freeze_bn_delay=500000):
+                                       quant_delay=0,
+                                       freeze_bn_delay=int(2e5)):
   """Rewrites a training input_graph in place for simulated quantization.
 
   This function has additional experimental options not (yet) available to
@@ -141,6 +149,14 @@ def experimental_create_training_graph(input_graph=None,
   the expected behavior of previously held references to nodes and tensors may
   change.
 
+  The default value of quant_delay is suitable for finetuning an already trained
+  floating point model (recommended).
+  If one wants to train a quantized model from scratch, quant_delay should be
+  set to the number of steps it take the floating point model to converge.
+  Quantization will be activated at this point and effectively finetune the
+  model. If quant_delay is not provided when training from scratch, training can
+  often fail.
+
   Args:
     input_graph: The tf.Graph to be transformed,if None then defaults to the
       default graph.