import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets, layers, optimizers, Sequential, metrics

def preprocess(x,y):
    x=tf.cast(x, dtype=tf.float32)/255.
    y=tf.cast(y, dtype=tf.int32)
    return x,y

# 加载数据集
(x,y),(x_test,y_test)=datasets.fashion_mnist.load_data()
print(x.shape,y.shape)

batchsz=128

# 数据集划分成batch
db=tf.data.Dataset.from_tensor_slices((x,y))
db=db.map(preprocess).shuffle(10000).batch(batchsz)

db_test=tf.data.Dataset.from_tensor_slices((x_test,y_test))
db_test=db_test.map(preprocess).batch(batchsz)

# 测试查看
db_iter=iter(db)
sample=next(db_iter)
print('batch:',sample[0].shape,sample[1].shape)


# 构建前馈网络
model=Sequential([
    layers.Dense(256,activation=tf.nn.relu),
    layers.Dense(128,activation=tf.nn.relu),
    layers.Dense(64,activation=tf.nn.relu),
    layers.Dense(32,activation=tf.nn.relu),
    layers.Dense(10),
])

# build
model.build(input_shape=[None,28*28])
model.summary()

optimizer=optimizers.Adam(lr=1e-3)
total_correct=0
total_num=0

# 训练
  for epoch in range(30):
        for step ,(x,y) in enumerate(db):
            x=tf.reshape(x,[-1,28*28])
            with tf.GradientTape() as tape:
                logits=model(x)
                y_onehot=tf.one_hot(y,depth=10)
                loss_mse=tf.reduce_mean(tf.losses.MSE(y_onehot,logits))
                loss_ce=tf.losses.categorical_crossentropy(y_onehot,logits,from_logits=True)
                loss_ce=tf.reduce_mean(loss_ce)
            grads=tape.gradient(loss_ce,model.trainable_variables)
            # 用定义好的optimizer对trainable_variables更新
            optimizer.apply_gradients(zip(grads,model.trainable_variables))
            
            if step%100==0:
                print(epoch,step,'loss:',float(loss_ce),float(loss_mse))
        
        # test
        for x,y in db_test:
            # x:[b,28,28]=>[b,784]
            # y:[b]
            x=tf.reshape(x,[-1,28*28])
            # [b,10]
            logits=model(x)
            # logits=>prob,[b,10]
            prob=tf.nn.softmax(logits,axis=1)
            # [b,10]=>[b]
            pred=tf.argmax(prob,axis=1)
            # pred.dtype是int64
            pred=tf.cast(pred,dtype=tf.int32)
            # pred:[b]
            # y:[b]
            # correct:[b]
            correct=tf.equal(pred,y)
            correct=tf.reduce_sum(tf.cast(correct,dtype=tf.int32))
            total_correct+=int(correct)
            total_num += x.shape[0]

        acc=total_correct/total_num
        print(epoch,'test acc:',acc)