Internal change

PiperOrigin-RevId: 299931864

Author: pkch

morph_net/framework/batch_norm_source_op_handler_test.py

@@ -65,18 +65,6 @@ def setUp(self):
     self.beta_op_group = orm.OpGroup(
         self.beta_op_slice, omit_source_op_slices=[self.beta_op_slice])
 
-    self.mean_op = g.get_operation_by_name(
-        'conv1/BatchNorm/AssignMovingAvg/sub_1')
-    self.mean_op_slice = orm.OpSlice(self.mean_op, orm.Slice(0, 5))
-    self.mean_op_group = orm.OpGroup(
-        self.mean_op_slice, omit_source_op_slices=[self.mean_op_slice])
-
-    self.std_op = g.get_operation_by_name(
-        'conv1/BatchNorm/AssignMovingAvg_1/sub_1')
-    self.std_op_slice = orm.OpSlice(self.std_op, orm.Slice(0, 5))
-    self.std_op_group = orm.OpGroup(
-        self.std_op_slice, omit_source_op_slices=[self.std_op_slice])
-
     # Create custom mapping of OpSlice and OpGroup in manager.
     self.mock_op_reg_manager = mock.create_autospec(orm.OpRegularizerManager)
 
@@ -91,7 +79,7 @@ def get_op_group(op_slice):
     self.mock_op_reg_manager.is_source_op.return_value = False
     self.mock_op_reg_manager.ops = [
         self.batch_norm_op, self.conv_op, self.relu_op, self.gamma_op,
-        self.beta_op, self.mean_op, self.std_op]
+        self.beta_op]
 
   def testAssignGrouping_NoNeighborGroups(self):
     self.op_slice_dict = {
@@ -100,8 +88,6 @@ def testAssignGrouping_NoNeighborGroups(self):
         self.relu_op: [self.relu_op_slice],
         self.gamma_op: [self.gamma_op_slice],
         self.beta_op: [self.beta_op_slice],
-        self.mean_op: [self.mean_op_slice],
-        self.std_op: [self.std_op_slice],
     }
 
     # No neighbor ops have groups.
@@ -128,7 +114,7 @@ def testAssignGrouping_NoNeighborGroups(self):
 
     # Verify manager processes grouping for input and output ops.
     self.mock_op_reg_manager.process_ops.assert_called_once_with(
-        [self.relu_op, self.mean_op, self.std_op, self.conv_op, self.gamma_op,
+        [self.relu_op, self.conv_op, self.gamma_op,
          self.beta_op])
     self.mock_op_reg_manager.process_ops_last.assert_called_once_with(
         [self.batch_norm_op])
@@ -140,8 +126,6 @@ def testAssignGrouping_ProcessNeighborGroups(self):
         self.relu_op: [self.relu_op_slice],
         self.gamma_op: [self.gamma_op_slice],
         self.beta_op: [self.beta_op_slice],
-        self.mean_op: [self.mean_op_slice],
-        self.std_op: [self.std_op_slice],
     }
 
     # All ops have groups.
@@ -151,8 +135,6 @@ def testAssignGrouping_ProcessNeighborGroups(self):
         self.relu_op_slice: self.relu_op_group,
         self.gamma_op_slice: self.gamma_op_group,
         self.beta_op_slice: self.beta_op_group,
-        self.mean_op_slice: self.mean_op_group,
-        self.std_op_slice: self.std_op_group,
     }
 
     # Call handler to assign grouping.
@@ -167,8 +149,7 @@ def testAssignGrouping_ProcessNeighborGroups(self):
 
     # Verify manager groups batch norm with outputs and inputs.
     self.mock_op_reg_manager.group_op_slices.assert_has_calls(
-        [mock.call([self.batch_norm_op_slice, self.relu_op_slice,
-                    self.mean_op_slice, self.std_op_slice]),
+        [mock.call([self.batch_norm_op_slice, self.relu_op_slice]),
          mock.call([self.batch_norm_op_slice, self.conv_op_slice,
                     self.gamma_op_slice, self.beta_op_slice])])
 
@@ -208,28 +189,12 @@ def testAssignGrouping_ProcessNeighborGroupsWithSlices(self):
     beta_op_group2 = orm.OpGroup(
         beta_op_slice_2_3, omit_source_op_slices=[beta_op_slice_2_3])
 
-    mean_op_slice_0_2 = orm.OpSlice(self.mean_op, orm.Slice(0, 2))
-    mean_op_slice_2_3 = orm.OpSlice(self.mean_op, orm.Slice(2, 1))
-    mean_op_group1 = orm.OpGroup(
-        mean_op_slice_0_2, omit_source_op_slices=[mean_op_slice_0_2])
-    mean_op_group2 = orm.OpGroup(
-        mean_op_slice_2_3, omit_source_op_slices=[mean_op_slice_2_3])
-
-    std_op_slice_0_2 = orm.OpSlice(self.std_op, orm.Slice(0, 2))
-    std_op_slice_2_3 = orm.OpSlice(self.std_op, orm.Slice(2, 1))
-    std_op_group1 = orm.OpGroup(
-        std_op_slice_0_2, omit_source_op_slices=[std_op_slice_0_2])
-    std_op_group2 = orm.OpGroup(
-        std_op_slice_2_3, omit_source_op_slices=[std_op_slice_2_3])
-
     self.op_slice_dict = {
         self.batch_norm_op: [batch_norm_op_slice_0_2, batch_norm_op_slice_2_3],
         self.conv_op: [conv_op_slice_0_2, conv_op_slice_2_3],
         self.relu_op: [relu_op_slice_0_2, relu_op_slice_2_3],
         self.gamma_op: [gamma_op_slice_0_2, gamma_op_slice_2_3],
         self.beta_op: [beta_op_slice_0_2, beta_op_slice_2_3],
-        self.mean_op: [mean_op_slice_0_2, mean_op_slice_2_3],
-        self.std_op: [std_op_slice_0_2, std_op_slice_2_3],
     }
 
     # All OpSlice have groups.
@@ -244,10 +209,6 @@ def testAssignGrouping_ProcessNeighborGroupsWithSlices(self):
         gamma_op_slice_2_3: gamma_op_group2,
         beta_op_slice_0_2: beta_op_group1,
         beta_op_slice_2_3: beta_op_group2,
-        mean_op_slice_0_2: mean_op_group1,
-        mean_op_slice_2_3: mean_op_group2,
-        std_op_slice_0_2: std_op_group1,
-        std_op_slice_2_3: std_op_group2,
     }
 
     # Call handler to assign grouping.
@@ -262,12 +223,10 @@ def testAssignGrouping_ProcessNeighborGroupsWithSlices(self):
 
     # Verify manager groups batch norm with outputs and inputs by slice.
     self.mock_op_reg_manager.group_op_slices.assert_has_calls(
-        [mock.call([batch_norm_op_slice_0_2, relu_op_slice_0_2,
-                    mean_op_slice_0_2, std_op_slice_0_2]),
+        [mock.call([batch_norm_op_slice_0_2, relu_op_slice_0_2]),
          mock.call([batch_norm_op_slice_0_2, conv_op_slice_0_2,
                     gamma_op_slice_0_2, beta_op_slice_0_2]),
-         mock.call([batch_norm_op_slice_2_3, relu_op_slice_2_3,
-                    mean_op_slice_2_3, std_op_slice_2_3]),
+         mock.call([batch_norm_op_slice_2_3, relu_op_slice_2_3]),
          mock.call([batch_norm_op_slice_2_3, conv_op_slice_2_3,
                     gamma_op_slice_2_3, beta_op_slice_2_3])])
 
@@ -282,8 +241,6 @@ def testAssignGrouping_NeighborsHaveSameGroup(self):
         self.relu_op: [self.batch_norm_op_slice],
         self.gamma_op: [self.batch_norm_op_slice],
         self.beta_op: [self.batch_norm_op_slice],
-        self.mean_op: [self.batch_norm_op_slice],
-        self.std_op: [self.batch_norm_op_slice],
     }
 
     # All ops have the same group.
@@ -320,8 +277,6 @@ def testAssignGrouping_NeighborsHaveSameGroup_ReprocessSources(self):
         self.relu_op: [self.relu_op_slice],
         self.gamma_op: [self.gamma_op_slice],
         self.beta_op: [self.beta_op_slice],
-        self.mean_op: [self.mean_op_slice],
-        self.std_op: [self.std_op_slice],
     }
 
     self.op_group_dict = {
@@ -353,9 +308,6 @@ def is_source_op(op):
         [self.batch_norm_op_slice, self.conv_op_slice, self.gamma_op_slice,
          self.beta_op_slice])
 
-    # Verify manager adds ungrouped output ops to queue.
-    self.mock_op_reg_manager.process_ops.assert_called_once_with(
-        [self.mean_op, self.std_op])
     self.mock_op_reg_manager.process_ops_last.assert_not_called()
 
   def testCreateRegularizer(self):

morph_net/framework/op_handler_util_test.py

@@ -59,11 +59,6 @@ def setUp(self):
 
     self.gamma_op = g.get_operation_by_name('conv1/BatchNorm/gamma/read')
     self.beta_op = g.get_operation_by_name('conv1/BatchNorm/beta/read')
-    self.decay_op = g.get_operation_by_name('conv1/BatchNorm/Const')
-    self.mean_op = g.get_operation_by_name(
-        'conv1/BatchNorm/AssignMovingAvg/sub_1')
-    self.std_op = g.get_operation_by_name(
-        'conv1/BatchNorm/AssignMovingAvg_1/sub_1')
 
     self.relu_op = g.get_operation_by_name('conv1/Relu')
     self.relu_op_slice = orm.OpSlice(self.relu_op, None)
@@ -113,8 +108,8 @@ def is_passthrough(op):
     self.mock_op_reg_manager.get_op_group.side_effect = get_op_group
     self.mock_op_reg_manager.is_passthrough.side_effect = is_passthrough
     self.mock_op_reg_manager.ops = [
-        self.batch_norm_op, self.gamma_op, self.beta_op, self.decay_op,
-        self.mean_op, self.std_op, self.conv_op, self.relu_op,
+        self.batch_norm_op, self.gamma_op, self.beta_op,
+        self.conv_op, self.relu_op,
         self.relu2_op, self.relu3_op, self.relu4_op, self.unfused_batch_norm_op,
         self.concat_op]
 
@@ -127,11 +122,10 @@ def testGetInputOps(self):
     input_ops = op_handler_util.get_input_ops(self.batch_norm_op,
                                               self.mock_op_reg_manager)
     self.assertEqual(expected_inputs, input_ops)
-    self.assertNotIn(self.decay_op, input_ops)
 
   def testGetOutputOps(self):
     # For batch norm, the expected outputs are mean, std, and ReLU.
-    expected_outputs = [self.relu_op, self.mean_op, self.std_op]
+    expected_outputs = [self.relu_op]
 
     # Check for expected output ops.
     self.assertEqual(
@@ -146,7 +140,6 @@ def testGetOpsWithoutGroups(self):
         self.conv_op: [self.conv_op_slice],
         self.gamma_op: [orm.OpSlice(self.gamma_op, None)],
         self.beta_op: [orm.OpSlice(self.beta_op, None)],
-        self.decay_op: [orm.OpSlice(self.decay_op, None)],
     }
 
     # Only batch norm and conv ops have groups.
@@ -155,21 +148,19 @@ def testGetOpsWithoutGroups(self):
         self.conv_op_slice: self.conv_op_group
     }
 
-    all_ops = [self.batch_norm_op, self.conv_op, self.gamma_op, self.beta_op,
-               self.decay_op]
+    all_ops = [self.batch_norm_op, self.conv_op, self.gamma_op, self.beta_op]
     # Batch norm and conv ops have groups.  The other ops do not have groups.
-    expected_ops = [self.gamma_op, self.beta_op, self.decay_op]
+    expected_ops = [self.gamma_op, self.beta_op]
     self.assertEqual(
         expected_ops,
         op_handler_util.get_ops_without_groups(
             all_ops, self.mock_op_reg_manager))
 
   def testRemoveNonPassthroughOps(self):
-    self._passthrough_ops = (self.gamma_op, self.decay_op, self.std_op)
+    self._passthrough_ops = (self.gamma_op,)
 
-    all_ops = [self.batch_norm_op, self.conv_op, self.gamma_op, self.beta_op,
-               self.decay_op, self.mean_op]
-    expected_ops = [self.gamma_op, self.decay_op]
+    all_ops = [self.batch_norm_op, self.conv_op, self.gamma_op, self.beta_op]
+    expected_ops = [self.gamma_op]
 
     self.assertListEqual(
         expected_ops,