在 TensorFlow 之中自定義訓(xùn)練

在前面兩節(jié)的學(xué)習(xí)之中，我們學(xué)習(xí)了如何進(jìn)行自定義微分操作，又學(xué)習(xí)了如何自定義模型，那么接下來(lái)這一小節(jié)我們便來(lái)學(xué)習(xí)如何進(jìn)行自定義的最后一步 —— 自定義訓(xùn)練。

在之前的學(xué)習(xí)之中，當(dāng)我們進(jìn)行訓(xùn)練的時(shí)候，我們采用的都是 fit () 函數(shù)。雖然說(shuō) fit () 函數(shù)給我們提供了很多的選項(xiàng)，但是如果想要更加深入的定制我們的尋來(lái)你過(guò)程，那么我們便需要自己編寫(xiě)訓(xùn)練循環(huán)。

在這節(jié)課之中，我們會(huì)采用一個(gè)對(duì) mnist 數(shù)據(jù)集進(jìn)行分類(lèi)的簡(jiǎn)單模型作為一個(gè)簡(jiǎn)單的示例來(lái)進(jìn)行演示，以此來(lái)幫助大家理解自定義訓(xùn)練的過(guò)程。因此該課程主要可以分為以下幾個(gè)部分：

• 自定義模型；
• 編寫(xiě)自定義循環(huán)；
• 如何在自定義循環(huán)中進(jìn)行模型的優(yōu)化工作。

1. 數(shù)據(jù)的準(zhǔn)備工作

首先我們要準(zhǔn)備好相應(yīng)的 mnist 數(shù)據(jù)集，這里采用往常的處理方式：使用內(nèi)置的 API 來(lái)獲取數(shù)據(jù)集：

import tensorflow as tf

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
# 數(shù)據(jù)歸一化
train_images = train_images / 255.0
test_images = test_images / 255.0


train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)

valid_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
valid_dataset = valid_dataset.batch(64)

在這里，不僅僅構(gòu)建了數(shù)據(jù)集，我們同樣對(duì)圖片數(shù)據(jù)進(jìn)行了歸一化的操作。同時(shí)，我們對(duì)數(shù)據(jù)進(jìn)行了分批次的處理，批次的大小維 64 。于此同時(shí)，我們對(duì)訓(xùn)練數(shù)據(jù)進(jìn)行了亂序處理。

2. 自定義模型

由于進(jìn)行 Mnist 圖像分類(lèi)的任務(wù)比較簡(jiǎn)單，因此我們可以定義一個(gè)較為簡(jiǎn)單的模型，這里的模型的結(jié)構(gòu)包含四層：

• Flattern 層：對(duì)二維數(shù)據(jù)進(jìn)行展開(kāi)；
• 第一個(gè) Dense 層：包含 128 個(gè)神經(jīng)元；
• 第二個(gè) Dense 層：包含 64 個(gè)神經(jīng)元；
• 最后一個(gè) Dense 分類(lèi)層；包含 10 個(gè)神經(jīng)元，對(duì)應(yīng)于我們的十個(gè)分類(lèi)。

class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.l1 = tf.keras.layers.Flatten()
        self.l2 = tf.keras.layers.Dense(128, activation='relu')
        self.l3 = tf.keras.layers.Dense(64, activation='relu')
        self.l4 = tf.keras.layers.Dense(10, activation='softmax')

    def call(self, inputs, training=True):
        x = self.l1(inputs)
        x = self.l2(x)
        x = self.l3(x)
        y = self.l4(x)
        return y
model = MyModel()

3. 定義訓(xùn)練循環(huán)

在做好準(zhǔn)備工作之后，我們便來(lái)到了我們的最重要的部分，也就是如何進(jìn)行自定義循環(huán)的構(gòu)建。

在自定義循環(huán)之前，我們要先做好準(zhǔn)備工作，分為如下幾步：

• 自定義損失函數(shù)：在大多數(shù)情況之下，內(nèi)置的損失函數(shù)以及足夠我們使用，比如交叉熵等損失函數(shù)；
• 自定義優(yōu)化器：優(yōu)化器決定了按照如何的策略進(jìn)行優(yōu)化，我們最常用的優(yōu)化器就是 Adam ，因此這里我們使用內(nèi)置的 Adam 優(yōu)化器；
• （可選）定義變量監(jiān)視器：用于監(jiān)視我們的訓(xùn)練過(guò)程的各種參數(shù)，在這里我們只使用一個(gè)來(lái)監(jiān)視我們的驗(yàn)證集合上的效果。

因此我們的代碼可以如下所示：

# 損失函數(shù)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
# 優(yōu)化器
optimizer = tf.keras.optimizers.Adam()
# 監(jiān)控驗(yàn)證機(jī)上的準(zhǔn)確率
val_acc = tf.keras.metrics.SparseCategoricalAccuracy()

然后我們便可以構(gòu)建自定義循環(huán)，自定義循環(huán)大致分為以下幾步：

• 編寫(xiě)一個(gè)循環(huán) Epoch 次的循環(huán)，Epoch 為訓(xùn)練的循環(huán)數(shù)；
• 在循環(huán)內(nèi)部對(duì)于數(shù)據(jù)集讀取每一個(gè) Batch，因?yàn)檫@里的 train_dataset 為可枚舉的，因此我們直接使用枚舉即可獲得每一個(gè)批次的訓(xùn)練樣本；
• 定義 tf.GradientTape () 梯度帶；
• 在梯度帶內(nèi)進(jìn)行模型的輸出，以及損失的求取；
• 在梯度帶外使用梯度帶求得模型所有參數(shù)的梯度，在這里我們可以使用 model.trainable_weights 來(lái)獲取所有可訓(xùn)練的參數(shù)；
• 使用優(yōu)化器按照求得的梯度對(duì)模型的參數(shù)進(jìn)行優(yōu)化，這里直接使用 optimizer.apply_gradients 函數(shù)即可完成優(yōu)化；
• （可選）進(jìn)行 Log 處理，打印出日志便于我們查看；
• （可選）在每個(gè) Epoch 的訓(xùn)練集的訓(xùn)練結(jié)束后，我們可以在測(cè)試集上查看結(jié)果，這里我們只查看了準(zhǔn)確率。

epochs = 3
for epoch in range(epochs):
    print("Start Training epoch " + str(epoch))
    
    # 取出每一個(gè)批次的數(shù)據(jù)
    for batch_i, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        # 在梯度帶內(nèi)進(jìn)行操作
        with tf.GradientTape() as tape:
          outputs = model(x_batch_train, training=True)
          loss_value = loss_fn(y_batch_train, outputs)

        # 求取梯度
        grads = tape.gradient(loss_value, model.trainable_weights)
        # 使用Optimizer進(jìn)行優(yōu)化
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

        # Log
        if batch_i % 100 == 0:
            print("Loss at batch %d: %.4f" % (batch_i, float(loss_value)))

    # 在驗(yàn)證集合上測(cè)試
    for batch_i, (x_batch_train, y_batch_train) in enumerate(valid_dataset):
        outputs = model(x_batch_train, training=False)
        # 更新追蹤器的狀態(tài)
        val_acc.update_state(y_batch_train, outputs)
    print("Validation acc: %.4f" % (float(val_acc.result()),))

    # 重置追蹤器
    val_acc.reset_states()

最終，我們可以得到如下結(jié)果：

Start Training epoch 0
Loss at batch 0: 0.1494
Loss at batch 100: 0.2155
Loss at batch 200: 0.1080
Loss at batch 300: 0.0231
Loss at batch 400: 0.1955
Loss at batch 500: 0.2019
Loss at batch 600: 0.0567
Loss at batch 700: 0.1099
Loss at batch 800: 0.0714
Loss at batch 900: 0.0364
Validation acc: 0.9691
Start Training epoch 1
Loss at batch 0: 0.0702
Loss at batch 100: 0.0615
Loss at batch 200: 0.0208
Loss at batch 300: 0.0158
Loss at batch 400: 0.0304
Loss at batch 500: 0.1193
Loss at batch 600: 0.0130
Loss at batch 700: 0.1353
Loss at batch 800: 0.1300
Loss at batch 900: 0.0056
Validation acc: 0.9715
Start Training epoch 2
Loss at batch 0: 0.0714
Loss at batch 100: 0.0066
Loss at batch 200: 0.0177
Loss at batch 300: 0.0086
Loss at batch 400: 0.0099
Loss at batch 500: 0.1621
Loss at batch 600: 0.1103
Loss at batch 700: 0.0049
Loss at batch 800: 0.0139
Loss at batch 900: 0.0111
Validation acc: 0.9754

大家可以發(fā)現(xiàn)，我們的模型在測(cè)試集合上達(dá)到了 97.54% 的準(zhǔn)確率。

同時(shí)我們可以發(fā)現(xiàn)，其實(shí)在第一個(gè) Epoch 之后，模型已經(jīng)達(dá)到了很好地效果，這是因?yàn)槲覀兊娜蝿?wù)比較簡(jiǎn)單，而且我們的模型擬合能力比較強(qiáng)。

4. 小結(jié)

這節(jié)課我們采用案例驅(qū)動(dòng)的方式，在對(duì)圖片進(jìn)行分類(lèi)的過(guò)程之中學(xué)會(huì)了如何自定義循環(huán)，以及如何在自定義循環(huán)的時(shí)候進(jìn)行優(yōu)化。