import tensorflow as tf from tensorflow import keras from tensorflow.keras import datasets import os # do not print irrelevant information # os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # x: [60k,28,28], [10,28,28] # y: [60k], [10k] (x, y), (x_test, y_test) = datasets.mnist.load_data() # transform Tensor # x: [0~255] ==》 [0~1.] x = tf.convert_to_tensor(x, dtype=tf.float32) / 255. y = tf.convert_to_tensor(y, dtype=tf.int32) x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) / 255. y_test = tf.convert_to_tensor(y_test, dtype=tf.int32)
# batch of 128 train_db = tf.data.Dataset.from_tensor_slices((x, y)).batch(128) test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(128) train_iter = iter(train_db) sample = next(train_iter)
# [b,784] ==> [b,256] ==> [b,128] ==> [b,10] # [dim_in,dim_out],[dim_out] w1 = tf.Variable(tf.random.truncated_normal([784, 256], stddev=0.1)) b1 = tf.Variable(tf.zeros([256])) w2 = tf.Variable(tf.random.truncated_normal([256, 128], stddev=0.1)) b2 = tf.Variable(tf.zeros([128])) w3 = tf.Variable(tf.random.truncated_normal([128, 10], stddev=0.1)) b3 = tf.Variable(tf.zeros([10]))
# learning rate lr = 1e-3
for epoch in range(10): # iterate db for 10 # tranin every train_db for step, (x, y) in enumerate(train_db): # x: [128,28,28] # y: [128] # [b,28,28] ==> [b,28*28] x = tf.reshape(x, [-1, 28 * 28]) with tf.GradientTape( ) as tape: # only data types of tf.variable are logged # x: [b,28*28] # h1 = x@w1 + b1 # [b,784]@[784,256]+[256] ==> [b,256] + [256] ==> [b,256] + [b,256] h1 = x @ w1 + tf.broadcast_to(b1, [x.shape[0], 256]) h1 = tf.nn.relu(h1) # [b,256] ==> [b,128] # h2 = x@w2 + b2 # b2 can broadcast automatic h2 = h1 @ w2 + b2 h2 = tf.nn.relu(h2) # [b,128] ==> [b,10] out = h2 @ w3 + b3 # compute loss # out: [b,10] # y:[b] ==> [b,10] y_onehot = tf.one_hot(y, depth=10) # mse = mean(sum(y-out)^2) # [b,10] loss = tf.square(y_onehot - out) # mean:scalar loss = tf.reduce_mean(loss) # compute gradients grads = tape.gradient(loss, [w1, b1, w2, b2, w3, b3]) # w1 = w1 - lr * w1_grad # w1 = w1 - lr * grads[0] # not in situ update # in situ update w1.assign_sub(lr * grads[0]) b1.assign_sub(lr * grads[1]) w2.assign_sub(lr * grads[2]) b2.assign_sub(lr * grads[3]) w3.assign_sub(lr * grads[4]) b3.assign_sub(lr * grads[5]) if step % 100 == 0: print(f'epoch:{epoch}, step: {step}, loss:{float(loss)}') # [w1,b1,w2,b2,w3,b3] total_correct, total_num = 0, 0 for step, (x, y) in enumerate(test_db): # [b,28,28] ==> [b,28*28] x = tf.reshape(x, [-1, 28 * 28]) # [b,784] ==> [b,256] ==> [b,128] ==> [b,10] h1 = tf.nn.relu(x @ w1 + b1) h2 = tf.nn.relu(h1 @ w2 + b2) out = h2 @ w3 + b3 # out: [b,10] ~ R # prob: [b,10] ~ (0,1) prob = tf.nn.softmax(out, axis=1) # [b,10] ==> [b] pred = tf.argmax(prob, axis=1) pred = tf.cast(pred, dtype=tf.int32) # y: [b] # [b], int32 correct = tf.cast(tf.equal(pred, y), dtype=tf.int32) correct = tf.reduce_sum(correct) total_correct += int(correct) total_num += x.shape[0] acc = total_correct / total_num print(f'test acc: {acc}')