import tensorflow as tf

def preporocess(x,y):
    x = tf.cast(x,dtype=tf.float32) / 255
    x = tf.reshape(x,(-1,28 *28))                   #   鋪平
    x = tf.squeeze(x,axis=0)
    # print('裡面x.shape:',x.shape)
    y = tf.cast(y,dtype=tf.int32)
    return x,y

def main():
    #   載入手寫數字資料
    mnist = tf.keras.datasets.mnist
    (train_x, train_y), (test_x, test_y) 
 
= mnist.load_data()
    #   處理資料
        #   訓練資料
    db = tf.data.Dataset.from_tensor_slices((train_x,train_y))    #   將x,y分成一一對應的元組
    db = db.map(preporocess)                                    #   執行預處理函式
    db = db.shuffle(60000).batch(2000)                          #   打亂加分組
        #   測試資料
    db_test = tf.data.Dataset.from_tensor_slices((test_x,test_y))
    db_test 
 
= db_test.map(preporocess)
    db_test = db_test.shuffle(10000).batch(10000)
    #   設定超參
    iter_num = 2000                                             #   迭代次數
    lr = 0.01                                                   #   學習率
    #   定義模型器和優化器
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(
 
256,activation='relu'),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dense(10)
    ])
    # model.build(input_shape=[None,28*28])                     #   事先檢視網路結構
    # model.summary()
    #   優化器
    # optimizer = tf.keras.optimizers.SGD(learning_rate=lr)
    optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

    #   迭代訓練
    db_iter = iter(db)
    for i in range(iter_num):
        for step,(x,y) in enumerate(db):
            with tf.GradientTape() as tape:
                logits = model(x)
                y_onehot = tf.one_hot(y,depth=10)
                # loss = tf.reduce_mean(tf.losses.MSE(y_onehot,logits))                                         #   差平方損失
                loss = tf.reduce_mean(tf.losses.categorical_crossentropy(y_onehot,logits,from_logits=True))     #   交叉熵損失
            grads = tape.gradient(loss,model.trainable_variables)                                               #   梯度
            grads,_ = tf.clip_by_global_norm(grads,15)                                                          #   梯度限幅
            optimizer.apply_gradients(zip(grads,model.trainable_variables))                                     #   更新引數
            if step % 10 == 0:
                pass
                # print('i:{} , step:{} , loss:{} '.format(i,step,loss))
        #   計算測試集準確率
        for (x,y) in db_test:
            logits = model(x)
            out = tf.nn.softmax(logits,axis=1)
            pre = tf.argmax(out,axis=1)
            pre = tf.cast(pre,dtype=tf.int32)
            print(pre.shape,y.shape)
            acc  = tf.equal(pre,y)
            acc = tf.cast(acc,dtype=tf.int32)
            acc = tf.reduce_mean(tf.cast(acc,dtype=tf.float32))
            print('i:{}'.format(i))
            print('acc:{}'.format(acc))

if __name__ == '__main__':
    main()