Tensorflow2.0 自定义训练

Tensorflow的自定义的训练可以提供用户更高的灵活度。此处以Minst训练进行讲解。

• 首先加载TF, 从datasets 中引入minst

import tensorflow as tf
minst=tf.keras.datasets.mnist
(x_train,y_train),(x_test,y_test)=minst.load_data()

• 建立模型，引入models的Sequential，此处使用一个BP神经网络

model=tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=[28,28]),
    tf.keras.layers.Dense(128,activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10,activation='softmax')
])

• 建立优化器和损失函数
• 优化器使用ADAM
• 损失函数使用交叉熵
• tf.GradientTape 自动计算梯度
• to_categorical 对y进行one-hot
• tape.gradient 计算梯度，必须在with之外，参数是损失函数和可训练权重。
• 使用优化器应用梯度，需要zip进行对应处理。

optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001)
cce=tf.keras.losses.CategoricalCrossentropy()
batch_size=64
class_num=10
epoch_num=10
for epoch in range(epoch_num):
    index=0
    for _ in range(int(len(y_train)/batch_size)):
        with tf.GradientTape() as tape:
            x=x_train[index:index+batch_size]
            y=y_train[index:index+batch_size]
            index+=batch_size
            y_pre=model(x)
            y=tf.keras.utils.to_categorical(y,class_num)
            losses=cce(y,y_pre)
        grads = tape.gradient(losses,model.trainable_weights)
        optimizer.apply_gradients(zip(grads,model.trainable_weights))

• 验证

model.compile(optimizer='adam',
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy'])
#训练模型
model.evaluate(x_test,y_test)