fix test keras layers

Signed-off-by: zehao-intel <[email protected]>

test/3x/tensorflow/keras/test_layers.py

@@ -1,4 +1,3 @@
-!/usr/bin/env python
 # -*- coding: utf-8 -*-
 #
 # Copyright (c) 2024 Intel Corporation
@@ -30,7 +29,51 @@
 logger = Logger().get_logger()
 
 
-def build_model():
+def build_model1():
+    # Load MNIST dataset
+    mnist = keras.datasets.mnist
+
+    # 60000 images in train and 10000 images in test, but we don't need so much for ut
+    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()
+    train_images, train_labels = train_images[:1000], train_labels[:1000]
+    test_images, test_labels = test_images[:200], test_labels[:200]
+
+    # Normalize the input image so that each pixel value is between 0 to 1.
+    train_images = train_images / 255.0
+    test_images = test_images / 255.0
+
+    # Define the model architecture.
+    model = keras.Sequential(
+        [
+            keras.layers.InputLayer(input_shape=(28, 28)),
+            keras.layers.Reshape(target_shape=(28, 28, 1)),
+            keras.layers.DepthwiseConv2D(3, 3, activation='relu', name="conv2d"),
+            keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            keras.layers.Flatten(),
+            keras.layers.Dense(10, name="dense"),
+        ]
+    )
+    # Train the digit classification model
+    model.compile(
+        optimizer="adam", loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"]
+    )
+
+    model.fit(
+        train_images,
+        train_labels,
+        epochs=1,
+        validation_split=0.1,
+    )
+
+    _, baseline_model_accuracy = model.evaluate(test_images, test_labels, verbose=0)
+
+    print("Baseline test accuracy:", baseline_model_accuracy)
+    if version1_gte_version2(tf.__version__, "2.16.1"):
+        model.save("baseline_model1.keras")
+    else:
+        model.save("baseline_model1")
+
+def build_model2():
     # Load MNIST dataset
     mnist = keras.datasets.mnist
 
@@ -128,7 +171,7 @@ def tearDownClass(self):
             shutil.rmtree(self.fp32_model_path, ignore_errors=True)
         os.environ["ITEX_ONEDNN_GRAPH"] = "0"
 
-    def test_static_quant_from_dict_default(self):
+    def test_depthwise_conv2d(self):
         logger.info("test_static_quant_from_dict_default")
         from neural_compressor.tensorflow import quantize_model
         from neural_compressor.tensorflow.keras import get_default_static_quant_config