# -- 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'