Tensorflow2.0基础-笔记-自定义训练

import tensorflow as tf

#这段是因为我笔记本显存不够加的
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)





#导入训练数据和测试数据
(train_image,train_labels),(test_images,test_labels)=tf.keras.datasets.mnist.load_data()

#对图片进行预处理，扩增一个维（记录卷积的通道数），图片数据归一化（x/255），转换数据类型
train_image=tf.expand_dims(train_image,-1)
train_image=tf.cast(train_image/255,tf.float32)
train_labels=tf.cast(train_labels,tf.int64)

test_images=tf.expand_dims(test_images,-1)
test_images=tf.cast(test_images/255,tf.float32)
test_labels=tf.cast(test_labels,tf.int64)






#使用slice切片导入
dataset=tf.data.Dataset.from_tensor_slices((train_image,train_labels))
test_dataset=tf.data.Dataset.from_tensor_slices((test_images,test_labels))






#打乱数据，设置一个batch为32个数据
dataset=dataset.shuffle(10000).batch(32) #没有设置repeat
test_dataset=test_dataset.batch(32)





#建立模型，2层卷积，最后打平成10个神经元
model=tf.keras.Sequential([
    tf.keras.layers.Conv2D(16,[3,3],activation='relu',input_shape=(28,28,1)),#input_shape=(None,None,1)
    tf.keras.layers.Conv2D(32,[3,3],activation='relu'),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(10) #不使用激活函数
    
])


#设置优化方法为Adam，loss函数为多元交叉熵
optimizer=tf.keras.optimizers.Adam()
loss_func=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) 
#from_logits在没有最后的激活函数的时候设置为True







#建立计算一个epoch的loss和accuracy的函数，train和test的数据都需要计算
train_loss=tf.keras.metrics.Mean('train_loss') #计算loss的均值
train_acc=tf.keras.metrics.SparseCategoricalAccuracy('train_accuracy')

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






#设置训练一个batch的train数据过程，（反向传播，更新梯度和变量）
def train_step(model,images,labels): 
    #tf.GradientTape()记录运算过程的变量梯度
    with tf.GradientTape() as t：
        pred=model(images)    #计算一个batch的预测值
        loss_step=loss_func(labels,pred)   #计算一个batch的损失函数
    grads=t.gradient(loss_step,model.trainable_variables)  #更新变量的梯度
    #利用跟新的梯度计算变量，model.trainable_variables表示网络中所有的变量
    optimizer.apply_gradients(zip(grads,model.trainable_variables))  
    train_loss(loss_step)  #传入计算一个epoch累计loss函数中
    train_acc(labels,pred) #传入计算一个epoch累计accuracy函数中

#设置训练一个batch的test数据过程，不需要反向传播
def test_step(model,images,labels): 
    pred=model(images)
    loss_step=loss_func(labels,pred)
    test_loss(loss_step)
    test_acc(labels,pred)







#设置总训练过程函数，epoch_time为运行epoch的次数
def train(epoch_time):
    for epoch in range(epoch_time):
    #在每一个epoch中都需要计算所有的train数据一次。

        for (batch,(images,labels)) in enumerate(dataset):
            train_step(model,images,labels)
        print('Epoch {} is finished，loss={:.3f}，acc={:.3f}'.format(epoch,train_loss.result(),train_acc.result()))
        
        for (batch,(images,labels)) in enumerate(test_dataset):
            test_step(model,images,labels)
        print('     test_loss={:.3f}，test_accuracy={:.3f}'.format(test_loss.result(),test_acc.result()))
        
        #一个epoch训练完后，重新初始化loss，accuracy的值
        train_loss.reset_states()
        train_acc.reset_states()
        test_loss.reset_states()
        test_acc.reset_states()



#训练10个epoch
train(10)