# 导入相应的工具包

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (7,7) # Make the figures a bit bigger
import tensorflow as tf
# 数据集

from tensorflow.keras.datasets import mnist
# 构建序列模型

from tensorflow.keras.models import Sequential
# 导入需要的层

from tensorflow.keras.layers import Dense, Dropout, Activation,BatchNormalization
# 导入辅助工具包

from tensorflow.keras import utils
# 正则化

from tensorflow.keras import regularizers

# 类别总数

nb_classes = 10
# 加载数据集

(X_train, y_train), (X_test, y_test) = mnist.load_data()
# 打印输出数据集的维度

print("训练样本初始维度", X_train.shape)
print("训练样本目标值初始维度", y_train.shape)

# 数据展示：将数据集的前九个数据集进行展示

for i in range(9):
    plt.subplot(3,3,i+1)
    # 以灰度图显示，不进行插值

    plt.imshow(X_train[i], cmap='gray', interpolation='none')
    # 设置图片的标题：对应的类别

    plt.title("数字{}".format(y_train[i]))

# 调整数据维度：每一个数字转换成一个向量

X_train = X_train.reshape(60000, 784)
X_test = X_test.reshape(10000, 784)
# 格式转换

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
# 归一化

X_train /= 255
X_test /= 255
# 维度调整后的结果

print("训练集：", X_train.shape)
print("测试集：", X_test.shape)

# 将目标值转换为热编码的形式

Y_train = utils.to_categorical(y_train, nb_classes)
Y_test = utils.to_categorical(y_test, nb_classes)

# 利用序列模型来构建模型

model = Sequential()
# 全连接层，共512个神经元，输入维度大小为784

model.add(Dense(512, input_shape=(784,)))
# 激活函数使用relu

model.add(Activation('relu')) 
# 使用正则化方法drouout 

model.add(Dropout(0.2))  
# 全连接层，共512个神经元,并加入L2正则化

model.add(Dense(512,kernel_regularizer=regularizers.l2(0.001)))
# BN层

model.add(BatchNormalization())
# 激活函数

model.add(Activation('relu'))
model.add(Dropout(0.2))
# 全连接层，输出层共10个神经元

model.add(Dense(10))
# softmax将神经网络输出的score转换为概率值

model.add(Activation('softmax')) 

# 模型编译，指明损失函数和优化器，评估指标

model.compile(loss='categorical_crossentropy', optimizer='adam',metrics=['accuracy'])

# batch_size是每次送入模型中样本个数，epochs是所有样本的迭代次数，并指明验证数据集

history = model.fit(X_train, Y_train,
          batch_size=128, epochs=4,verbose=1,
          validation_data=(X_test, Y_test))

# 绘制损失函数的变化曲线

plt.figure()
# 训练集损失函数变换

plt.plot(history.history["loss"], label="train_loss")
# 验证集损失函数变化

plt.plot(history.history["val_loss"], label="val_loss")
plt.legend()
plt.grid()


# 绘制准确率的变化曲线

plt.figure()
# 训练集准确率

plt.plot(history.history["accuracy"], label="train_acc")
# 验证集准确率

plt.plot(history.history["val_accuracy"], label="val_acc")
plt.legend()
plt.grid()

# 添加tensoboard观察

tensorboard = tf.keras.callbacks.TensorBoard(log_dir='./graph', histogram_freq=1,
                                                write_graph=True,write_images=True)

# 训练

history = model.fit(X_train, Y_train,
          batch_size=128, epochs=4,verbose=1,callbacks=[tensorboard],
          validation_data=(X_test, Y_test))

# 指定存在文件的目录，打开下面命令

tensorboard --logdir="./"

# 模型测试

score = model.evaluate(X_test, Y_test, verbose=1)
# 打印结果

print('测试集准确率:', score)

# 保存模型架构与权重在h5文件中

model.save('my_model.h5')
# 加载模型：包括架构和对应的权重

model = tf.keras.models.load_model('my_model.h5')