Removed potentially accumulated errors from testSeparableConvQuantize_

PiperOrigin-RevId: 379835554

Author: Xhark

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_transforms_test.py

@@ -290,9 +290,8 @@ def testSeparableConv1DQuantize_(self, kwargs):
   @parameterized.named_parameters(
       ('padding_valid', {'padding': 'valid'}),
       ('padding_same', {'padding': 'same'}),
-      # TODO(b/186666265): tighten the tolerance to 1e-5.
       ('padding_same_dilation_2',
-       {'padding': 'same', 'dilation_rate': 2}, 0.19),
+       {'padding': 'same', 'dilation_rate': 2}),
       ('strides', {'strides': 2}),
       ('dilation_rate', {'dilation_rate': 2}),
       ('depth_multiplier', {'depth_multiplier': 2}),
@@ -307,7 +306,7 @@ def testSeparableConv1DQuantize_(self, kwargs):
           'pointwise_constraint': tf.keras.constraints.min_max_norm(0., 2.),
           'bias_constraint': tf.keras.constraints.unit_norm()})
   )
-  def testSeparableConvQuantize_(self, kwargs, tolerance=1e-5):
+  def testSeparableConvQuantize_(self, kwargs):
     kwargs['filters'] = 2
     kwargs['kernel_size'] = 3
     num_samples = 2
@@ -338,17 +337,20 @@ def testSeparableConvQuantize_(self, kwargs, tolerance=1e-5):
 
     # Ensure model is equivalent, and training results are the same.
     sepconv_model.compile(loss='categorical_crossentropy', optimizer='sgd')
-    sepconv_model.fit(x, y, epochs=100)
     transformed_model.compile(loss='categorical_crossentropy', optimizer='sgd')
-    transformed_model.fit(x, y, epochs=100)
 
-    # Over a long training cycle with constraints and regularizers, the model
-    # can build very minute differences.
-    self.assertAllClose(
-        sepconv_model.predict(x),
-        transformed_model.predict(x),
-        atol=tolerance,
-        rtol=tolerance)
+    epochs = 100
+    for _ in range(epochs):
+      sepconv_model.fit(x, y, epochs=1, verbose=2)
+      transformed_model.fit(x, y, epochs=1, verbose=2)
+      self.assertAllClose(
+          sepconv_model.get_weights(),
+          transformed_model.get_weights())
+      # To prevent accumulated numerical errors.
+      transformed_model.set_weights(sepconv_model.get_weights())
+      self.assertAllClose(
+          sepconv_model.predict(x),
+          transformed_model.predict(x))
 
   # TODO(pulkitb): Add individual tests for the following transforms.
   # Conv2DReshapeBatchNormQuantize, Conv2DReshapeBatchNormReLUQuantize