Tensorflow 2.0 mnist

# -- coding: utf-8 --
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.utils.data_utils import get_file
from tensorflow.python.util.tf_export import keras_export
from tensorflow import keras
from tensorflow.keras import layers, datasets, Sequential, optimizers, metrics
from tensorflow_core.python.eager import profiler
import datetime as dt

base_path = '/root'
mnist_path = base_path + '/mnist.npz'
@keras_export('keras.datasets.mnist.load_data')
def load_data(path):
    with np.load(path) as f:
      x_train, y_train = f['x_train'], f['y_train']
      x_test, y_test = f['x_test'], f['y_test']
      return (x_train, y_train), (x_test, y_test)
(x_train, y_train), (x_test, y_test) = load_data(mnist_path)

# 数据预处理
def train_preprocess(x_train, y_train):
    x_train = tf.cast(x = x_train, dtype = tf.float32) / 255.
    y_train = tf.cast(x = y_train, dtype = tf.int32)
    y_train = tf.one_hot(indices = y_train, depth = 10)
    return x_train, y_train

def test_preprocess(x_test, y_test):
    x_test = tf.cast(x = x_test, dtype = tf.float32) / 255.
    y_test = tf.cast(x = y_test, dtype = tf.int32)
    return x_test, y_test

train_db = tf.data.Dataset.from_tensor_slices(tensors=(x_train, y_train))
train_db = train_db.map(map_func=train_preprocess).shuffle(buffer_size=1000).batch(batch_size=128)

test_db = tf.data.Dataset.from_tensor_slices(tensors=(x_test, y_test))
test_db = test_db.map(map_func=test_preprocess).batch(batch_size=128)

# 建立网络模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(units=512, activation=tf.nn.relu),
    tf.keras.layers.Dense(units=256, activation=tf.nn.relu),
    tf.keras.layers.Dense(units=128, activation=tf.nn.relu),
    tf.keras.layers.Dense(units=32, activation=tf.nn.relu),
    tf.keras.layers.Dense(units=10),
])
model.build(input_shape=[None, 28 * 28])
model.summary()
optimizer = tf.keras.optimizers.Adam(learning_rate = 1e-4)

def main():
    for epoch in range(100):
        for step, (x_train, y_train) in enumerate(train_db):
            profiler.start()
            x_train = tf.reshape(tensor = x_train, shape = [-1, 28 * 28])

            with tf.GradientTape() as tape:
                logits = model(x_train)
                loss = tf.losses.categorical_crossentropy(y_true = y_train, y_pred = logits, from_logits = True)
                loss = tf.reduce_mean(input_tensor = loss)

            gradient = tape.gradient(target = loss, sources = model.trainable_variables)
            optimizer.apply_gradients(zip(gradient, model.trainable_variables))
            
            profiler_result = profiler.stop()
            if step % 100 == 0:
                log_dir= base_path + "/profile/" + dt.datetime.now().strftime("%Y%m%d-%H%M%S")
                profiler.save(log_dir, profiler_result)
                print('{}epoch, {} step, loss{}'.format(epoch, step, float(loss)))

        total_correct = 0
        total_num = 0
        for step, (x_test, y_test) in enumerate(test_db):
            x_test = tf.reshape(tensor = x_test, shape = [-1, 28 * 28])
            logits = model(x_test)
            probability = tf.nn.softmax(logits = logits, axis = 1)
            prediction = tf.argmax(input = probability, axis = 1)
            prediction = tf.cast(x = prediction, dtype = tf.int32)
            correct = tf.equal(x = prediction, y = y_test)
            correct = tf.cast(x = correct, dtype = tf.int32)
            correct = tf.reduce_sum(input_tensor = correct)
    
            total_correct += int(correct)
            total_num += x_test.shape[0]

        accuracy = total_correct / total_num
        print('accuracy:',accuracy)
if __name__ == '__main__':
    main()

　　其中，mnist.npz数据可以在此链接下载'https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz'