tf2的自定义训练
简单来说,什么是自定义训练?我们使用过tensoeflow2都知道直接model.compile进行编译后,model.fit就可以直接进行训练。这样确实很方便,但是其实在某种程度限制了自身的灵活性。下面本文便以Mnist数据集为例,完成自定义项目的流程。
项目流程
第一步,读入数据,将数据维度转换标准格式
(train_images, train_labels), (test_images, test_labels) = keras.datasets.mnist.load_data()
train_images, test_images = train_images / 255.0, test_images / 255.0
train_images.shape
train_images = np.expand_dims(train_images, axis=3) #输入需要四维
test_images = np.expand_dims(test_images, axis=3)
train_images.shape
第二步,将数据做成dataset对象
BATCH_SIZE = 128
train_image_ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_count = train_images.shape[0]
train_image_ds = train_image_ds.shuffle(train_count).batch(BATCH_SIZE)
test_image_ds = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
test_count = test_images.shape[0]
第三步,使用搭积木的方法构建神经网络
model = keras.Sequential([
tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3), padding='same',input_shape=(28, 28, 1), activation='relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), padding='same',activation='relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Flatten(), #拉平
tf.keras.layers.Dense(128, activation='relu'), #第一个全连接层
tf.keras.layers.Dense(10, activation='softmax') #第二个全连接层,由于是输出层,所以使用softmax做激活函数
])
model.summary()
第四步,构建评价指标和优化器
epoch_loss_avg = tf.keras.metrics.Mean('train_loss') #一个batch的平均损失函数值
train_accuracy = tf.keras.metrics.Accuracy() #一个batch的准确率
epoch_loss_avg_test = tf.keras.metrics.Mean('test_loss')
test_accuracy = tf.keras.metrics.Accuracy()
optimizer = tf.keras.optimizers.Adam() #构建优化器
第五步,设计训练和测试函数
在设计函数时候,我们采用tensorflow2自动求导的功能进行求解。特别地,测试函数无需自动求导。
def train_step(model, images, labels):
with tf.GradientTape() as t:
pred = model(images)
pred = tf.argmax(model(images),axis=1) #拿到可能性最大的类别
pred = tf.cast(pred,dtype=tf.float32) #转换数据类别
labels = tf.cast(labels,dtype=tf.float32)
loss_step = tf.keras.losses.CategoricalCrossentropy()(labels, pred) #损失函数
grads = t.gradient(loss_step, model.trainable_variables) #求导
optimizer.apply_gradients(zip(grads, model.trainable_variables)) #优化梯度
epoch_loss_avg(loss_step) ##一个batch的平均损失函数值
train_accuracy(labels, tf.cast(pred>0, tf.int32))
第六步,构建训练过程
train_loss_results = []
train_acc_results = []
test_loss_results = []
test_acc_results = []
def test_step(model, images, labels):
pred = model(images, training=False)
pred = tf.argmax(model(images),axis=1)
pred = tf.cast(pred,dtype=tf.float32)
labels = tf.cast(labels,dtype=tf.float32)
loss_step = tf.keras.losses.CategoricalCrossentropy()(labels, pred)
epoch_loss_avg_test(loss_step)
test_accuracy(labels, tf.cast(pred>0, tf.int32))
num_epochs = 5
for epoch in range(num_epochs):
for imgs_, labels_ in train_image_ds:
#labels_ = tf.one_hot(labels_,depth=10)
train_step(model, imgs_, labels_) #训练
print('.', end='')
print()
train_loss_results.append(epoch_loss_avg.result()) #记录,方便可视化
train_acc_results.append(train_accuracy.result())
for imgs_, labels_ in test_image_ds:
labels_ = tf.one_hot(labels_,depth=10)
test_step(model, imgs_, labels_) #测试
test_loss_results.append(epoch_loss_avg_test.result())
test_acc_results.append(test_accuracy.result())
print('Epoch:{}: loss: {:.3f}, accuracy: {:.3f}, test_loss: {:.3f}, test_accuracy: {:.3f}'.format(
epoch + 1,
epoch_loss_avg.result(),
train_accuracy.result(),
epoch_loss_avg_test.result(),
test_accuracy.result()
)) #打印
epoch_loss_avg.reset_states() #归0,统计下一个batch的
train_accuracy.reset_states()
epoch_loss_avg_test.reset_states()
test_accuracy.reset_states()