[Tensorflow]從Pytorch到TF2的學習之路 - Custom Model & Custom training

【[Tensorflow]從Pytorch到TF2的學習之路】所有文章:

前言

TF2.0中採用了和Pytorch相同的Eager Mode，並且使用了大量的Keras API，使得我們可以像Pytorch一樣透過建立動態圖來操作我們的運算的同時，能夠更有效率地來設計&訓練我們的深度學習模型。但在寫法上，Pytorch和TF2的到底有著什麼樣的差異呢?

這篇文章針對客製化模型和自定義訓練循環這兩個部分來撰寫，分別介紹使用Pytorch和TF2的不同，期望讓大家能夠快速地切換到不同的深度學習框架上。

客製化模型的不同之處

不囉So，咱們直接上code，先看一下兩個版本的程式碼，然後一併介紹他們的差異

首先先看Pytorch版本

class CustomModel(torch.nn.Module):
    def __init__(self):
        super(self).__init__()
        # define your layer
        self.fc1 = torch.nn.Linear(...)

    def forward(self, x):
        output = self.fc1(x)
        return output

再來看TF2版本

class CustomModel(tf.keras.Model):
    def __init__(self):
        super(self).__init__()
        # define your layer
        self.fc1 = tf.keras.layers.Dense(...)

    def call(self, x):
        output = self.fc1(x)
        return output

這邊簡單介紹一下差異的部分:

Pytorch
1. 在Pytorch中，如果要自定義一個模型，必須要繼承torch.nn.Module，然後實作__init__(self)和forward(self)兩個function
2. layer的定義是透過torch.nn來宣告
TF2
1. 而在TF2中，則是透過了Keras Model API來實作，因此必須繼承tf.keras.Model並且實作__init__(self)和call(self)兩個function
  - 實際上tf.keras.Model是繼承自tf.Module，所以也可以直接繼承tf.Module來實作自己的Model
  - 不過繼承tf.keras.Model還可以使用Model.fit(), Model.evaluate()和Model.save()這些操作，比方說官方文件的範例code，繼承tf.Module的話就要自己寫訓練循環
2. layer的定義是透過tf.keras.layers來宣告
3. tf.Module以及被他繼承的類別皆提供很方便的properties: variables and trainable_variables，使得我們能夠很好的管理模型的參數

最後，在一篇教學文章中看到了下面一個問題: 为什么模型类是重载 call() 方法而不是 __call__() 方法？，覺得這個觀念也值得一提，所以把內容直接複製上來

在 Python 中，对类的实例 myClass 进行形如 myClass() 的调用等价于 myClass.__call__() （具体请见本章初 “前置知识” 的 __call__() 部分）。那么看起来，为了使用 ypred = model(X) 的形式调用模型类，应该重写 \_call_() 方法才对呀？原因是 Keras 在模型调用的前后还需要有一些自己的内部操作，所以暴露出一个专门用于重载的 call() 方法。 tf.keras.Model 这一父类已经包含 \_call__() 的定义。 __call__() 中主要调用了 call() 方法，同时还需要在进行一些 keras 的内部操作。这里，我们通过继承 tf.keras.Model 并重载 call() 方法，即可在保持 keras 结构的同时加入模型调用的代码。

自定義訓練循環的不同之處

再來來看一下一個訓練的起手式寫法上，兩者有著什麼樣的差異:

首先從Pytorch看起，在[Pytorch]zero_grad()和backward()使用技巧中，我們提到了Pytorch訓練的起手式:

for idx, (batch_x, batch_y) in enumerate(data_loader):
    output = model(batch_x)
    loss = criterion(output, batch_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

短短幾行裡面做了很多事情:

將data傳入model進行forward propagation
計算loss
清空前一次的gradient
根據loss進行back propagation，計算gradient
做gradient descent

接下來我們來看一下TF2.0是如何撰寫的，儘管繼承tf.keras.Model的Model可以無腦model.fit()和model.predict()，但我們先看如果要自己寫訓練循環時應該怎麼做:

for batch_idx in range(batch_num):
    # 由於我還不太熟Pytorch的dataloader在TF2要怎麼寫，所以這裡寫得跟上面有點不同
    # 主要就是取出資料的意思
    batch_X, batch_y = data_loader.get_batch(batch_size)
    with tf.GradientTape() as tape:
        output = model(batch_X)
        loss = criterion(batch_y, output) # calculate loss
    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grad_and_vars=zip(grads, model.variables))

TF2做了以下的事情:

透過tf.GradientTape()紀錄並建構正向傳播的計算圖被包覆的操作
將data傳入model進行forward propagation
計算loss
根據loss計算對model.variables的梯度(使用tf.GradientTape())
做gradient descent更新梯度
- 透過tf.keras.optimizer更新，並且apply_gradients需要接收grad_and_vars參數(需要將gradient和參數用zip包起來餵進去)

對於tf.GradientTape()，我覺得知乎上面這篇文章: tensorflow计算图与自动求导——tf.GradientTape寫得很詳細，以下節錄幾個重點:

預設watch_accessed_variables=True，也就是會記錄所有可訓練變數(trainable=True的tf.Variable)
也可以用tape.watch()來觀測特定的變數，即使他trainable=False
預設persistent = False，也就是和Pytorch一樣，計算圖一但執行完就會銷毀(為了節省記憶體)，如果需要保留計算圖可以設置persistent = True
在用tape.gradient()計算梯度時，如果計算失敗(例如計算圖中兩個點根本沒有連接)，預設會返回None，但其實也可以指定參數unconnected_gradients=tf.UnconnectedGradients.ZERO來設置返回0

到這裡，了解了模型的寫法和起手式的寫法後，原本對於Pytorch較熟悉的人應該就能較輕鬆的寫出一個可以跑的TF2程式(也希望對於TF2的使用者這篇能夠幫助你快速熟悉Pytorch的寫法)。

前言

客製化模型的不同之處

自定義訓練循環的不同之處

References