文 / 軟體工程實習生 Xuechen Li
來源：TensorFlow 公眾號
Eager Execution 可簡化 TensorFlow 中的模型構建體驗，而 Graph Execution 可提供優化，以加快模型執行速度及提高儲存效率。本篇博文展示瞭如何編寫 TensorFlow 程式碼，以便將藉助 tf.keras API 並使用 Eager Execution 構建的模型轉換為圖表，最終藉助 tf.estimator API 的支援，在 Cloud TPU 上部署此模型。
注：tf.keras 連結
www.tensorflow.org/guide/keras
tf.estimator 連結
www.tensorflow.org/guide/estim…
我們使用可逆殘差網路（RevNet、Gomez 等）作為示例。接下來的部分假設讀者對卷積神經網路和 TensorFlow 有基本瞭解。您可以在此處找到本文的完整程式碼（為確保程式碼在所有設定中正常執行，強烈建議您使用 tf-nightly 或 tf-nightly-gpu）。
RevNet
RevNet 與殘差網路（ResNet、He 等）類似，只不過他們是可逆的，在給定輸出的情況下可重建中間計算。此技術的好處之一是我們可以通過重建啟用來節省記憶體，而不是在訓練期間將其全部儲存在記憶體中（回想一下，由於鏈式法則有此要求，因此我們需要中間結果來計算有關輸入的梯度）。相比傳統架構上的一般反向傳播，這使我們可以適應較大的批次大小，並可訓練更具深度的模型。具體來說，此技術的實現方式是通過使用一組巧妙構建的方程來定義網路：
其中頂部和底部方程組分別定義正演計算和其反演計算。這裡的 x1 和 x2 是輸入（從整體輸入 x 中拆分出來），y1 和 y2 是輸出，F 和 G 是 ConvNet。這使我們能夠在反向傳播期間精準重建啟用，如此一來，在訓練期間便無需再儲存這些資料。
使用 tf.keras.Model 定義正向和反向傳遞
假設我們使用 “ResidualInner” 類來例項化函式 F 和 G，我們可以通過子類化 tf.keras.Model 來定義可逆程式碼塊，並通過替換上面的方程中所示的 call 方法來定義正向傳遞：
1 class Residual(tf.keras.Model):
2 def __init__(self, filters):
3 super(Residual, self).__init__()
4 self.f = ResidualInner(filters=filters, strides=(1, 1))
5 self.g = ResidualInner(filters=filters, strides=(1, 1))
6
7 def call(self, x, training=True):
8 x1, x2 = tf.split(x, num_or_size_splits=2, axis=self.axis)
9 f_x2 = self.f(x2, training=training)
10 y1 = f_x2 + x1
11 g_y1 = self.g(y1, training=training)
12 y2 = g_y1 + x2
13 return tf.concat([y1, y2], axis=self.axis)
複製程式碼這裡的 training 引數用於確定批標準化的狀態。啟用 Eager Execution 後，批標準化的執行平均值會在 training=True 時自動更新。執行等效圖時，我們需要使用 get_updates_for 方法手動獲取批標準化更新。
要構建節省記憶體的反向傳遞，我們需要使用 tf.GradientTape 作為上下文管理器來跟蹤梯度（僅在有需要時）：
注：tf.GradientTape 連結
www.tensorflow.org/api_docs/py…
1 def backward_grads(self, y, dy, training=True):
2 dy1, dy2 = dy
3 y1, y2 = y
4
5 with tf.GradientTape() as gtape:
6 gtape.watch(y1)
7 gy1 = self.g(y1, training=training)
8 grads_combined = gtape.gradient(
9 gy1, [y1] + self.g.trainable_variables, output_gradients=dy2)
10 dg = grads_combined[1:]
11 dx1 = dy1 + grads_combined[0]
12 x2 = y2 - gy1
13
14 with tf.GradientTape() as ftape:
15 ftape.watch(x2)
16 fx2 = self.f(x2, training=training)
17 grads_combined = ftape.gradient(
18 fx2, [x2] + self.f.trainable_variables,output_gradients=dx1)
19 df = grads_combined[1:]
20 dx2 = dy2 + grads_combined[0]
21 x1 = y1 - fx2
22
23 x = x1, x2
24 dx = dx1, dx2
25 grads = df + dg
26
27 return x, dx, grads
複製程式碼您可以在論文的 “演算法 1” 中找到確切的一組梯度計算（我們在程式碼中簡化了使用變數 z1 的中間步驟）。此演算法經過精心設計，在給定輸出和有關輸出的損失梯度的情況下，我們可以在每個可逆程式碼塊內，計算有關輸入和模型變數的梯度及重建輸入。呼叫 tape.gradient(y, x)，即可計算有關 x 的 y 梯度。我們也可使用引數 output_gradients 來明確應用鏈式法則。
使用 Eager Execution 來加快原型設計速度
使用 Eager Execution 進行原型設計的一個明顯好處是採用命令式操作。我們可以立即獲得結果，而不用先構建圖表，然後再初始化要執行的會話。
例如，我們通過由一般反向傳播計算的梯度來比較可逆反向傳播梯度，從而驗證我們的模型：
1 block = Residual()
2 x = tf.random_normal(shape=(N, C, H, W))
3 dy = tf.random_normal(shape=(N, C, H, W))
4 with tf.GradientTape() as tape:
5 tape.watch(x)
6 y = block(x)
7 # Compute true grads
8 dx_true = tape.gradient(y, x, output_gradients=dy)
9
10 # Compute grads from reconstruction
11 dx, _ = block.backward_grads(x, y, dy)
12
13 # Check whether the difference is below a certain 14 threshold
thres = 1e-6
15 diff_abs = tf.reshape(abs(dx - dx_true), [-1])
16 assert all(diff_abs < thres)
複製程式碼在上面的片段中，dx_true 是一般反向傳播返回的梯度，而 dx 是執行可逆反向傳播後返回的梯度。Eager Execution 整合了原生 Python，如此一來，all 和 abs 等函式便可直接應用於 Tensor。
使用 tf.train.Checkpoint 儲存和載入檢查點
為確保能夠使用 Eager Execution 和 Graph Execution 儲存和載入檢查點，TensorFlow 團隊建議您使用 tf.train.Checkpoint API。
為了儲存模型，我們使用想要儲存的所有物件建立了一個 tf.train.Checkpoint 例項。這個例項可能包括我們的模型、我們使用的優化器、學習率安排和全域性步驟：
1 checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer,
2 learning_rate=learning_rate, global_step=global_step)
複製程式碼我們可以按照下面的方法儲存和還原特定的已訓練例項：
1 checkpoint.save(file_prefix)
2 checkpoint.restore(save_path)
複製程式碼使用 tf.contrib.eager.defun 提升 Eager Execution 效能
由於解讀 Python 程式碼會產生開銷，Eager Execution 有時會比執行等效圖要慢。通過使用 tf.contrib.eager.defun 將由 TensorFlow 運算組成的 Python 函式編譯成可呼叫的 TensorFlow 圖表，可以彌補這種效能差距。在訓練深度學習模型時，我們通常可以在三個主要位置應用 tf.contrib.eager.defun：
- 正演計算
- 梯度的反演計算
- 將梯度應用於變數
例如，我們可以按以下方式 defun 正向傳遞和梯度計算：
1 tfe = tf.contrib.eager
2 model.call = tfe.defun(model.call)
3 model.compute_gradients = tfe.defun(model.compute_gradients)
複製程式碼要 defun 優化器的應用梯度步驟，我們需要將其包裝在另一個函式內：
1 def apply_gradients(optimizer, gradients, variables, global_step=None):
2 optimizer.apply_gradients(
3 zip(gradients, variables), global_step=global_step)
4 apply_gradients = tfe.defun(apply_gradients)
複製程式碼tf.contrib.eager.defun 正處於積極開發中，將其加以應用是一項不斷髮展的技術。如需更多資訊，請檢視其文件字串。
注：其文件字串連結
github.com/tensorflow/…
使用 tf.contrib.eager.defun 包裝 Python 函式會使 TensorFlow API 在 Python 函式中進行呼叫，以構建圖表，而不是立即執行運算，從而優化整個程式。並非所有 Python 函式都可成功轉換為等效圖，特別是帶有動態控制流的函式（例如，Tensor contents 中的 if 或 while）。tf.contrib.autograph 是一種相關工具，可以增加能夠轉換為 TensorFlow 圖表的 Python 程式碼的表面積。截至 2018 年 8 月，使用 defun 整合 Autograph 的工作仍在進行中。
注：tf.contrib.autograph 連結
www.tensorflow.org/guide/autog…
使用 TFRecords 和 tf.data.Dataset 構建輸入管道
Eager Execution 與 tf.data.Dataset API 相容。我們可以讀取 TFRecords 檔案：
1 dataset = tf.data.TFRecordDataset(filename)
2 dataset = dataset.repeat(epochs).map(parser).batch(batch_size)
複製程式碼為提升效能，我們還可使用預取函式並調整 num_parallel_calls。
由於資料集由影像和標籤對組成，在 Eager Execution 中迴圈使用此資料集非常簡單。在本例中，我們甚至不需要明確定義迭代器：
1 for image, label in dataset:
2 logits = model(image, training=True)
3 ...
複製程式碼使用估算器包裝 Keras 模型並以圖表形式執行
由於 tf.keras API 也支援圖表構建，因此使用 Eager Execution 構建的相同模型也可用作提供給估算器的圖表構建函式，但程式碼稍有更改。要修改使用 Eager Execution 構建的 RevNet 示例，我們只需使用 model_fn 包裝 Keras 模型，並按照 tf.estimator API 的指示使用此模型。
1 def model_fn(features, labels, mode, params):
2 model = RevNet(params["hyperparameters"])
3 if mode == tf.estimator.ModeKeys.TRAIN:
4 optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)
5 logits, saved_hidden = model(features, training=True)
6 grads, loss = model.compute_gradients(saved_hidden, labels, training=True)
7 with tf.control_dependencies(model.get_updates_for(features)):
8 train_op = optimizer.apply_gradients(zip(grads, model.trainable_variables))
9 return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
複製程式碼您可以使用 the tf.data API 照常定義 tf.estimator API 所需的 input_fn，並從 TFRecords 中讀取資料。
使用 TPU Estimator 包裝 Keras 模型以進行 Cloud TPU 訓練
使用 Estimator 包裝模型和輸入管道使模型可以在 Cloud TPU 上執行。
所需步驟如下：
設定 Cloud TPU 的特定配置
從 tf.estimator.Estimator 切換到 tf.contrib.tpu.TPUEstimator
使用 tf.contrib.tpu.CrossShardOptimizer 包裝常用優化器
注：Cloud TPU 連結
github.com/tensorflow/…
配置連結
www.tensorflow.org/api_docs/py…
tf.contrib.tpu.TPUEstimator 連結
www.tensorflow.org/api_docs/py…
tf.contrib.tpu.CrossShardOptimizer 連結
www.tensorflow.org/api_docs/py…
如需瞭解具體說明，請檢視 RevNet 示例資料夾中的 TPU 估算器指令碼。我們希望日後可以使用 tf.contrib.tpu.keras_to_tpu_model 進一步簡化使 Keras 模型在 TPU 上執行的流程。
注：TPU 估算器指令碼連結
github.com/tensorflow/…
tf.contrib.tpu.keras_to_tpu_model 連結
github.com/tensorflow/…
可選：模型效能
與一般反向傳播相比，有了 tf.GradientTape，再加上無需額外正向傳遞的梯度計算可簡化流程，我們能夠僅以 25% 的計算開銷來執行 RevNet 的可逆反向傳播。
圖中藍色和橙色的曲線分別表示隨著全域性步驟的增加，一般反向傳播和可逆反向傳播的每秒取樣率。該圖來自在單個 Tesla P100 上使用批次大小為 32 的模擬 ImageNet 資料訓練的 RevNet-104。
為了驗證所節省的記憶體，我們在訓練過程中繪製記憶體使用情況。藍色和黑色曲線分別是一般和可逆反向傳播。該圖記錄了使用批次大小為 128 的模擬 ImageNet 資料訓練 RevNet-104 圖表模式的 100 次迭代。該圖是在 CPU 上進行訓練時由 mprof 生成，以便我們使用一般反向傳播以相同的批次大小進行訓練。
結論
我們以 RevNet 為例，展示瞭如何使用 Eager Execution 和 tf.keras API 對機器學習模型快速進行原型設計。這不僅可簡化模型構建體驗，而且我們輕易就能將模型轉換為估算器，並在 Cloud TPU 上進行部署，以獲得高效能。您可以在此處找到本文的完整程式碼。此外，請務必檢視使用 Eager Execution 的其他示例。
注：此處連線誒
github.com/tensorflow/…
其他示例連結
github.com/tensorflow/…