import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import models
from tensorflow.keras import optimizers
from tensorflow.keras import datasets
from tensorflow.python.keras.utils import np_utils


(x_train,y_train),(x_test,y_test)=datasets.mnist.load_data()
print(x_train.shape,x_test.shape)
print(y_train.shape,y_test.shape)


x_train=x_train.reshape(x_train.shape[0],x_train.shape[1],x_train.shape[2],1)
x_test=x_test.reshape(x_test.shape[0],x_test.shape[1],x_test.shape[2],1)
x_train=x_train/255
x_test=x_test/255
y_train=np_utils.to_categorical(y_train)
y_test=np_utils.to_categorical(y_test)
print(x_train.shape,x_test.shape)
print(y_train.shape,y_test.shape)


inp_img=layers.Input(shape=(28,28,1))
layer1_conv=layers.Conv2D(32,(3,3),padding='same',activation='relu')(inp_img)
layer1_pool=layers.AveragePooling2D((2,2))(layer1_conv) 
layer2_conv=layers.Conv2D(64,(3,3),padding='same',activation='relu')(layer1_pool)
layer2_pool=layers.AveragePooling2D((2,2))(layer2_conv)  
layer3_conv=layers.Conv2D(128,(3,3),padding='same',activation='relu')(layer2_pool)
layer3_pool=layers.AveragePooling2D((7,7))(layer3_conv)   
layer4_fc=layers.Flatten()(layer3_pool)
pred=layers.Dense(10,activation='softmax')(layer4_fc)
model=models.Model(inp_img,pred)
adam=optimizers.Adam(lr=0.01)
model.compile(optimizer=adam,loss='categorical_crossentropy',metrics=['acc'])
model.summary()


model.fit(x_train,y=y_train,epochs=10,batch_size=32)
model.save('mnist_cnn.h5')
loss,acc=model.evaluate(x_test,y_test)
print('\ntest loss: ',loss)
print('\ntest accuracy: ',acc)