【ATU Book-i.MX8系列 - TFLite】Tensorflow 2. x 建立手写识别

#---------------------------------------------------------------------------------------------------------------
# 读取手写识别数据
#---------------------------------------------------------------------------------------------------------------
import tensorflow as tf
import tensorflow.keras.layers as layers
import numpy as np 
import os
from MNIST_data import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot = True)
print('\nLoading MNIST DataBase\n',
      'train number=',mnist.train.num_examples,"\n", 'vakidation number=',mnist.validation.num_examples,"\n",
      'test number=',mnist.test.num_examples,"\n", 'train image shape', mnist.train.images.shape,"\n",
      'label image shape', mnist.train.labels.shape,"\n", 'train image size' , mnist.train.images[0].size,"\n")
​

***  Tensorflow 2.x 版本已合并 Keras API ，故以下所介绍的范例部分会用到 Keras 的代码。 

#---------------------------------------------------------------------------------------------------------------
# 建立 卷积神经网路(动态图)
#---------------------------------------------------------------------------------------------------------------
inputs_ = layers.Input(shape=(28, 28, 1)) #输入层 input layer
C1_Conv = layers.Conv2D(filters=16,kernel_size=5, activation='relu')(inputs_)#卷积层 convolution layer 1
C1_Pool = layers.MaxPooling2D(2, 2)(C1_Conv) #池化层 pool layer 1
C2_Conv = layers.Conv2D(32, 5, activation = tf.nn.relu)(C1_Pool) #卷积层 convolution layer 2
C2_Pool = layers.MaxPooling2D(2, 2)(C2_Conv) #池化层 pool layer 2
flat1   = layers.Flatten()(C2_Pool) #平坦层 flatten layer （将多个维度的输入压平成一个维度）
Dense1  = layers.Dense((7*7*36), activation = tf.nn.relu)(flat1) #全部连阶层 fully connnected layer
Dropout = layers.Dropout(0.3)(Dense1)
outputs = layers.Dense(10, activation="softmax")(Dropout) #输出层 output layer
model   = tf.keras.Model(inputs=inputs_, outputs=outputs) 

#---------------------------------------------------------------------------------------------------------------
# 建立卷积神经网路(训练模型)
#---------------------------------------------------------------------------------------------------------------
#设定训练网路(函式、变数)
optimizer = tf.keras.optimizers.Adam() #最佳化设定 optimizer
loss_object = tf.keras.losses.SparseCategoricalCrossentropy() #损失函数
train_loss = tf.keras.metrics.Mean(name='train_loss') 
test_loss = tf.keras.metrics.Mean(name='test_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy') #准确度
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

@tf.function #启动 Auto Graph
def train_step(image, label):
    with tf.GradientTape() as tape:  
        predictions = model(image) #计算准确度
        loss = loss_object(label, predictions) #计算损失
    gradients = tape.gradient(loss, model.trainable_variables) #利用 GrudientTape 纪录梯度资讯
    optimizer.apply_gradients(zip(gradients, model.trainable_variables)) #梯度最佳化
    train_loss(loss)
    train_accuracy(label, predictions)
    
@tf.function #启动Auto Graph
def test_step(images, labels):
    predictions = model(images) #计算准确度
    t_loss = loss_object(labels, predictions) #计算损失
    test_loss(t_loss)
    test_accuracy(labels, predictions)

#读取资料库
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

#建构训练资料
x_train = x_train.astype(np.float32).reshape(-1, 28, 28) #调整资料大小
x_train = np.expand_dims(x_train, axis=3) #扩展资料
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(128) #拆分批次资料
#建构测试资料
x_test = x_test.astype(np.float32).reshape(-1, 28, 28) #调整资料大小
x_test = np.expand_dims(x_test, axis=3) #扩展资料
test_dataset  = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_dataset  = test_dataset.batch(128) #拆分批次资料

#开始训练模型 tranning model
epoch_list=[];accuracy_list=[];loss_list=[];
for epoch in range(20):
    for image, label in train_dataset:
        train_step(image, label)
        
    for test_images, test_labels in test_dataset:
        test_step(test_images, test_labels)    
        
    template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}\n'
    print(template.format(epoch+1, train_loss.result(), train_accuracy.result(),
                              test_loss.result() , test_accuracy.result()))
    
    epoch_list.append(epoch)
    loss_list.append(format(train_loss.result()))
    accuracy_list.append(format(train_accuracy.result()))

    # Reset the metrics for the next epoch
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()   

#---------------------------------------------------------------------------------------------------------------
# 预测结果
#---------------------------------------------------------------------------------------------------------------
import matplotlib.pyplot as plt

def plt_mnist(x, y, plt_num=5):
    plt.figure(figsize=(8,8))
    for index, (image, label) in enumerate(zip(x[0:plt_num], y[0:plt_num])):
        plt.subplot(2, 3, (index+1))
        plt.imshow(np.reshape(image, (28,28)), cmap=plt.cm.gray)
        plt.title('{}: label: {}\n'.format('Predict', label), fontsize = 10)

predict_ = model.predict(x_test)
predict = [tf.argmax(e).numpy() for e in predict_]
plt_mnist(x_test, predict)

tf.saved_model.save( model, "Tensorflow_2_HandWrite" )