在 TensorFlow 之中自定義訓練

在前面兩節的學習之中，我們學習了如何進行自定義微分操作，又學習了如何自定義模型，那么接下來這一小節我們便來學習如何進行自定義的最后一步 —— 自定義訓練。

在之前的學習之中，當我們進行訓練的時候，我們采用的都是 fit () 函數。雖然說 fit () 函數給我們提供了很多的選項，但是如果想要更加深入的定制我們的尋來你過程，那么我們便需要自己編寫訓練循環。

在這節課之中，我們會采用一個對 mnist 數據集進行分類的簡單模型作為一個簡單的示例來進行演示，以此來幫助大家理解自定義訓練的過程。因此該課程主要可以分為以下幾個部分：

• 自定義模型；
• 編寫自定義循環；
• 如何在自定義循環中進行模型的優化工作。

1. 數據的準備工作

首先我們要準備好相應的 mnist 數據集，這里采用往常的處理方式：使用內置的 API 來獲取數據集：

import tensorflow as tf

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
# 數據歸一化
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)

在這里，不僅僅構建了數據集，我們同樣對圖片數據進行了歸一化的操作。同時，我們對數據進行了分批次的處理，批次的大小維 64 。于此同時，我們對訓練數據進行了亂序處理。

2. 自定義模型

由于進行 Mnist 圖像分類的任務比較簡單，因此我們可以定義一個較為簡單的模型，這里的模型的結構包含四層：

• Flattern 層：對二維數據進行展開；
• 第一個 Dense 層：包含 128 個神經元；
• 第二個 Dense 層：包含 64 個神經元；
• 最后一個 Dense 分類層；包含 10 個神經元，對應于我們的十個分類。

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. 定義訓練循環

在做好準備工作之后，我們便來到了我們的最重要的部分，也就是如何進行自定義循環的構建。

在自定義循環之前，我們要先做好準備工作，分為如下幾步：

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

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

# 損失函數
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
# 優化器
optimizer = tf.keras.optimizers.Adam()
# 監控驗證機上的準確率
val_acc = tf.keras.metrics.SparseCategoricalAccuracy()

然后我們便可以構建自定義循環，自定義循環大致分為以下幾步：

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

epochs = 3
for epoch in range(epochs):
    print("Start Training epoch " + str(epoch))
    
    # 取出每一個批次的數據
    for batch_i, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        # 在梯度帶內進行操作
        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進行優化
        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)))

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

    # 重置追蹤器
    val_acc.reset_states()

最終，我們可以得到如下結果：

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

大家可以發現，我們的模型在測試集合上達到了 97.54% 的準確率。

同時我們可以發現，其實在第一個 Epoch 之后，模型已經達到了很好地效果，這是因為我們的任務比較簡單，而且我們的模型擬合能力比較強。

4. 小結

這節課我們采用案例驅動的方式，在對圖片進行分類的過程之中學會了如何自定義循環，以及如何在自定義循環的時候進行優化。