For LayerPattern, match config items with regex

PiperOrigin-RevId: 324901092

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer.py

@@ -104,19 +104,19 @@ def _match_pattern(self, target, pattern):
   def _match_layer(self, layer, pattern):
     """Check if specific layer matches the pattern."""
 
-    if self.candidate_layers and layer['config'][
-        'name'] not in self.candidate_layers:
+    if self.candidate_layers and \
+        layer['config']['name'] not in self.candidate_layers:
       return False
 
     if not self._match_pattern(layer['class_name'], pattern.class_name):
       return False
 
     layer_config = layer['config']
     for key, value in pattern.config.items():
-      # This comparison should probably use the serialized value.
-      # Consider adding regex support to key/values as well. This will allow
-      # negative matches as well.
-      if layer_config.get(key) != value:
+      # Either the provided value should equal the config value, or
+      # be a regex match to str(value).
+      if not (self._match_pattern(str(layer_config.get(key)), str(value)) or \
+              layer_config.get(key) == value):
         return False
 
     return True

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer_test.py

@@ -527,6 +527,16 @@ def matched(self):
     def reset(self):
       self._matched = False
 
+  @parameterized.parameters(['sequential', 'functional'])
+  def testPatternMatches_ConfigParamsRegex(self, model_type):
+    pattern = LayerPattern('Dense', config={'name': 'dense.*'})
+    transform = self.VerifyMatch(pattern)
+
+    model = self._simple_dense_model(model_type)
+
+    ModelTransformer(model, [transform]).transform()
+    self.assertTrue(transform.matched())
+
   @parameterized.parameters(['sequential', 'functional'])
   def testPatternShouldOnlyMatch_CandidateLayers(self, model_type):
     pattern = LayerPattern('ReLU', inputs=[LayerPattern('Dense')])
@@ -680,5 +690,6 @@ class MyModel(keras.Model):
     with self.assertRaises(ValueError):
       ModelTransformer(MyModel(), [self.ReplaceDenseLayer()]).transform()
 
+
 if __name__ == '__main__':
   tf.test.main()