import tensorflow as tf
import datetime

(train_image, train_labels), (test_image, test_labels)= tf.keras.datasets.mnist.load_data()
train_image = tf.expand_dims(train_image,-1)  # 原先形狀為(60000,28,28),現在變為(60000,28,28,1) 擴充套件維度，也就是通道維度
# 如果使用引數1，那麼形狀就變為了(60000,1,28,28)
test_image = tf.expand_dims(test_image,-1)
test_image = test_image / 255
train_image = train_image / 255  # 進行歸一化
train_image = tf.cast(train_image, tf.float32)  # 將資料型別轉換為float,因為只有float才能進行自動微分
test_image = tf.cast(test_image, tf.float32)
train_labels = tf.cast(train_labels,tf.int64)
test_labels = tf.cast(test_labels, tf.int64)
dataset = tf.data.Dataset.from_tensor_slices((train_image,train_labels))  # 將圖片和標籤進行對應組合，這個函式是將第一維進行拆分
# 也就是可以拆分為60000個單獨的資料
dataset = dataset.shuffle(1000).batch(32)  # 對資料進行混洗以及繫結32個為一組
test_dataset = tf.data.Dataset.from_tensor_slices((test_image,test_labels))
test_dataset = test_dataset.batch(32)

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, [3,3], activation='relu', input_shape=(None,None,1)),  # None表示只要是灰度圖都可以，沒有規定大小
    tf.keras.layers.Conv2D(32, [3,3], activation='relu'),
    tf.keras.layers.GlobalMaxPool2D(),
    tf.keras.layers.Dense(10,activation='softmax')  # 這裡沒有進行啟用函式，那麼就需要再後面的loss函式中進行一些操作
])

optimizer = tf.keras.optimizers.Adam()
loss_func = tf.keras.losses.SparseCategoricalCrossentropy()
def loss(model,x,y):
    y_ = model(x)
    return loss_func(y,y_)

train_loss = tf.keras.metrics.Mean('train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy('train_accuracy')
test_loss = tf.keras.metrics.Mean('test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy('test_accuracy')

def train_step(model,images,labels):
    with tf.GradientTape() as t:
        pred = model(images)
        loss_step = loss_func(labels,pred)
    grads = t.gradient(loss_step,model.trainable_variables)
    optimizer.apply_gradients(zip(grads,model.trainable_variables))
    train_loss(loss_step)
    train_accuracy(labels,pred)

def test_step(model, images, labels):
    pred = model(images)
    loss_step = loss_func(labels,pred)
    test_loss(loss_step)
    test_accuracy(labels,pred)

current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
train_log_dir = 'logs/gradient_tape' + current_time + "train"
test_log_dir = 'logs/gradient_tape' + current_time + "test"
train_writer = tf.summary.create_file_writer(train_log_dir)
test_writer = tf.summary.create_file_writer(test_log_dir)

def train(epoches):
    for epoch in range(epoches):
        for (batch,(images,labels)) in enumerate(dataset):
            train_step(model,images,labels)

        with train_writer.as_default():
            tf.summary.scalar('loss',train_loss.result(),step=epoch)
            tf.summary.scalar('acc',train_accuracy.result(),step=epoch)
        for (batch,(images,labels)) in enumerate(test_dataset):
            test_step(model,images,labels)
            # print('*',end='')
        with test_writer.as_default():
            tf.summary.scalar('test_loss', test_loss.result(), step=epoch)
            tf.summary.scalar('test_acc', test_accuracy.result(), step=epoch)
        template = 'Epoch {},Loss:{},Acc:{},Test Loss :{},Test Acc:{}'
        print(template.format(epoch+1,train_loss.result(),train_accuracy.result()*100,test_loss.result(),test_accuracy.result()*100))

        train_loss.reset_states()
        train_accuracy.reset_states()
        test_loss.reset_states()
        test_accuracy.reset_states()

train(10)