Better testing for batch normalization.

Author: LaurentMazare

src/graph/layer.ml

@@ -8,29 +8,28 @@ let batch_normalization ?(decay = 0.9) t ~update_moments ~dims ~feature_count =
   let beta = Var.create [ feature_count ] ~type_ ~init:zero in
   let gamma = Var.create [ feature_count ] ~type_ ~init:one in
   let batch_moments = Ops.moments t ~dims:(List.init dims ~f:Fn.id) in
-  let beta_with_update =
-    let update_beta =
-      (* EWMA update. *)
-      Ops.assignSub
-        beta
-        Ops.(one_minus_decay * (beta - batch_moments.mean))
-    in
-    Ops.identity ~control_inputs:[ Node.P update_beta ] beta
+  let beta_with_update ~control_inputs =
+    (* EWMA update. *)
+    Ops.assignSub
+      beta
+      Ops.(one_minus_decay * (beta - batch_moments.mean))
+      ~control_inputs
   in
-  let gamma_with_update =
-      (* EWMA update. *)
-    let update_gamma =
-      Ops.assignSub
-        gamma
-        Ops.(one_minus_decay * (gamma - batch_moments.variance))
-    in
-    Ops.identity ~control_inputs:[ Node.P update_gamma ] gamma
+  let gamma_with_update ~control_inputs =
+    (* EWMA update. *)
+    Ops.assignSub
+      gamma
+      Ops.(one_minus_decay * (gamma - batch_moments.variance))
+      ~control_inputs
   in
   let beta, gamma =
     match update_moments with
-    | `always -> beta_with_update, gamma_with_update
+    | `always ->
+      beta_with_update ~control_inputs:[], gamma_with_update ~control_inputs:[]
     | `not_in_testing testing ->
-      Ops.cond testing ~if_true:beta ~if_false:beta_with_update,
-      Ops.cond testing ~if_true:gamma ~if_false:gamma_with_update
+      let beta ~control_inputs:_ = beta in
+      let gamma ~control_inputs:_ = gamma in
+      Ops.cond_with_control_inputs testing ~if_true:beta ~if_false:beta_with_update,
+      Ops.cond_with_control_inputs testing ~if_true:gamma ~if_false:gamma_with_update
   in
   Ops.normalize t { mean = beta; variance = gamma }

src/graph/ops_manual.ml

@@ -18,6 +18,7 @@ let get_shape ?shape values =
 
 let const_float
     ?(name = "Const")
+    ?(control_inputs = [])
     ?shape
     ~type_
     values
@@ -28,7 +29,7 @@ let const_float
     ~op_name:(Op_name.of_string "Const")
     ~output_type:type_
     ~inputs:[]
-    ~control_inputs:[]
+    ~control_inputs
     ~attributes:[
       "dtype", Type (P type_);
       "value", Tensor_float { type_ = P type_; shape; values };
@@ -37,6 +38,7 @@ let const_float
 
 let const_int
     ?(name = "Const")
+    ?(control_inputs = [])
     ?shape
     ~type_
     values
@@ -47,7 +49,7 @@ let const_int
     ~op_name:(Op_name.of_string "Const")
     ~output_type:type_
     ~inputs:[]
-    ~control_inputs:[]
+    ~control_inputs
     ~attributes:[
       "dtype", Type (P type_);
       "value", Tensor_int { type_ = P type_; shape; values };
@@ -217,19 +219,24 @@ let normalize ?(epsilon = 1e-12) t { mean; variance } =
   let epsilon = scalar ~type_:(Node.output_type t) epsilon in
   Ops_generated.rsqrt (variance + epsilon) * (t - mean)
 
-let cond t ~if_true ~if_false =
+let cond_with_control_inputs t ~if_true ~if_false =
   let t_false, t_true = Ops_generated.switch t t in
   let if_true =
-    Ops_generated.identity if_true
+    if_true
       (* It is important to keep the [identity] below as control inputs do not handle
          ports. *)
       ~control_inputs:[ Node.P (Ops_generated.identity t_true) ]
   in
   let if_false =
-    Ops_generated.identity if_false
+    if_false
       (* It is important to keep the [identity] below as control inputs do not handle
          ports. *)
       ~control_inputs:[ Node.P (Ops_generated.identity t_false) ]
   in
   Ops_generated.merge [ if_true; if_false ]
   |> fst
+
+let cond t ~if_true ~if_false =
+  cond_with_control_inputs t
+    ~if_true:(fun ~control_inputs -> Ops_generated.identity ~control_inputs if_true)
+    ~if_false:(fun ~control_inputs -> Ops_generated.identity ~control_inputs if_false)

src/graph/ops_manual.mli

@@ -35,13 +35,15 @@ val cd : ?shape:int list -> float list -> [ `double ] Node.t
 (* Some more refined constant creation functions. *)
 val const_float
   :  ?name:string
+  -> ?control_inputs:Node.p list
   -> ?shape:int list
   -> type_:([< `float | `double ] as 'dtype) Node.Type.t
   -> float list
   -> 'dtype Node.t
 
 val const_int
   :  ?name:string
+  -> ?control_inputs:Node.p list
   -> ?shape:int list
   -> type_:([< `int32 | `int64 ] as 'dtype) Node.Type.t
   -> int list
@@ -155,6 +157,16 @@ val normalize
   -> 'a moments
   -> 'a Node.t
 
+(* If [if_true] and [if_false] use their [control_input] argument to build a node
+   this node will only be evaluated if necessary. *)
+val cond_with_control_inputs
+  :  [ `bool ] Node.t
+  -> if_true:(control_inputs:Node.p list -> 'a Node.t)
+  -> if_false:(control_inputs:Node.p list -> 'a Node.t)
+  -> 'a Node.t
+
+(* [if_true] and [if_false] will always be evaluated because of the 'not-so lazy'
+   behavior of TensorFlow switch. *)
 val cond
   :  [ `bool ] Node.t
   -> if_true:'a Node.t

tests/operator_tests.ml

@@ -92,39 +92,55 @@ let test_vector () =
 
 let test_batch_normalization () =
   let batch = Ops.placeholder ~type_:Double [ 3; 4 ] in
-  let testing = Ops.placeholder ~type_:Bool [] in
+  let testing = Ops.placeholder ~type_:Bool [ 1 ] in
   let ops =
     Layer.batch_normalization
       (Ops.Placeholder.to_node batch)
-      ~decay:0.
-      ~update_moments:`always
+      ~decay:0.5
+      ~update_moments:(`not_in_testing (Ops.Placeholder.to_node testing))
       ~dims:1
       ~feature_count:4
+    |> Ops.reduce_sum ~dims:[ 0 ]
   in
   let batch_tensor = Tensor.create2 Float64 3 4 in
   let testing_tensor = Tensor.create0 Int8_unsigned in
   Tensor.copy_elt_list testing_tensor [ 0 ];
-  let tensor =
-    Tensor.copy_elt_list batch_tensor
-      [ 0.; 4.; 0.; 8.
-      ; 0.; 4.; 9.; 8.
-      ; 0.; 4.; 3.; 3.
-      ];
-    Session.run Session.Output.(double ops)
-      ~inputs:
-        [ Session.Input.double batch batch_tensor
-        ; Session.Input.bool testing testing_tensor
-        ]
-  in
-  Tensor.print (Tensor.P tensor)
+  for i = 0 to 4 do
+    if i >= 3
+    then Tensor.copy_elt_list testing_tensor [ 1 ];
+    let tensor =
+      Tensor.copy_elt_list batch_tensor
+        [ 0.; 4.; 0.; 8.
+        ; 0.; 4.; 9.; 8.
+        ; 0.; 4.; 3.; 3.
+        ];
+      Session.run Session.Output.(double ops)
+        ~inputs:
+          [ Session.Input.double batch batch_tensor
+          ; Session.Input.bool testing testing_tensor
+          ]
+    in
+    let blessed_values =
+      if i = 0 then       [ 0.; 12.; 12.; 19. ]
+      else if i = 1 then [ 0.; 8.485281; 2.190890; 5.247275 ]
+      else if i = 2 then [ 0.; 6.000000; 0.914991; 2.260197 ]
+      else if i = 3 then [ 0.; 4.242641; 0.426401; 1.063611 ]
+      else if i = 4 then [ 0.; 4.242641; 0.426401; 1.063611 ]
+      else assert false
+    in
+    assert_vector tensor ~expected_value:blessed_values ~tol:1e-6
+  done
 
 let test_cond true_false =
   let testing = Ops.placeholder ~type_:Bool [ 1 ] in
   let true_false = if true_false then 1 else 0 in
+  let int32_with_control_inputs ~control_inputs v =
+    Ops.const_int ~shape:[] ~type_:Int32 ~control_inputs [ v ]
+  in
   let cond =
-    Ops.cond (Ops.Placeholder.to_node testing)
-      ~if_true:Ops.one32
-      ~if_false:Ops.zero32
+    Ops.cond_with_control_inputs (Ops.Placeholder.to_node testing)
+      ~if_true:(int32_with_control_inputs 1)
+      ~if_false:(int32_with_control_inputs 0)
   in
   let testing_tensor = Tensor.create0 Int8_unsigned in
   Tensor.copy_elt_list testing_tensor [ true_false ];
@@ -147,6 +163,6 @@ let test_cond true_false =
 let () =
   test_scalar ();
   test_vector ();
+  test_batch_normalization ();
   test_cond true;
-  test_cond false;
-  test_batch_normalization ()
+  test_cond false