Keras Model Creation using Model Subclass

Keras_Model_Subclass

#Import required libraries
import tensorflow as tf
from keras.datasets import mnist

#Creating the MNIST Dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test= x_train/255.0, x_test/255.0
x_train=x_train[...,tf.newaxis].astype('float32')
x_test=x_test[...,tf.newaxis].astype('float32')

#Building the model using Model sunc;ass
class Custom_Model(tf.keras.Model):
    def __init__(self):
        super(Custom_Model, self).__init__()
        self.conv1= tf.keras.layers.Conv2D(32,3, activation='relu')
        self.flatten= tf.keras.layers.Flatten()
        self.dense1= tf.keras.layers.Dense(128, activation='relu')
        self.dense2= tf.keras.layers.Dense(10, activation='softmax')
    def call(self, inputs):
        x= self.conv1(inputs)
        x= self.flatten(x)
        x= self.dense1(x)
        x= self.dense2(x)
        return x
model= Custom_Model()

# compile the mode;
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Fit the model on MNIST dataset
model.fit(x_train, y_train, epochs=10)

#Evaluate the model on the test data
test_loss, test_acc=model.evaluate(x_test, y_test)