Add decay/epsilon/updates_collections parameters to Inception v3 arg_scope

PiperOrigin-RevId: 182143473

1 files changed, 9 insertions, 3 deletions

diff --git a/tensorflow/contrib/slim/python/slim/nets/inception_v3.py b/tensorflow/contrib/slim/python/slim/nets/inception_v3.py
index a0e4b36058..afe261e43a 100644
--- a/tensorflow/contrib/slim/python/slim/nets/inception_v3.py
+++ b/tensorflow/contrib/slim/python/slim/nets/inception_v3.py
@@ -680,6 +680,9 @@ def _reduced_kernel_size_for_small_input(input_tensor, kernel_size):
 
 def inception_v3_arg_scope(weight_decay=0.00004,
                            batch_norm_var_collection='moving_vars',
+                           batch_norm_decay=0.9997,
+                           batch_norm_epsilon=0.001,
+                           updates_collections=ops.GraphKeys.UPDATE_OPS,
                            use_fused_batchnorm=True):
   """Defines the default InceptionV3 arg scope.
 
@@ -687,6 +690,9 @@ def inception_v3_arg_scope(weight_decay=0.00004,
     weight_decay: The weight decay to use for regularizing the model.
     batch_norm_var_collection: The name of the collection for the batch norm
       variables.
+    batch_norm_decay: Decay for batch norm moving average
+    batch_norm_epsilon: Small float added to variance to avoid division by zero
+    updates_collections: Collections for the update ops of the layer
     use_fused_batchnorm: Enable fused batchnorm.
 
   Returns:
@@ -694,11 +700,11 @@ def inception_v3_arg_scope(weight_decay=0.00004,
   """
   batch_norm_params = {
       # Decay for the moving averages.
-      'decay': 0.9997,
+      'decay': batch_norm_decay,
       # epsilon to prevent 0s in variance.
-      'epsilon': 0.001,
+      'epsilon': batch_norm_epsilon,
       # collection containing update_ops.
-      'updates_collections': ops.GraphKeys.UPDATE_OPS,
+      'updates_collections': updates_collections,
       # Use fused batch norm if possible.
       'fused': use_fused_batchnorm,
       # collection containing the moving mean and moving variance.