使用 Eager Execution 編碼並運行圖表：以經過 RevNet 優化代碼爲例

文 / 軟件工程實習生 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：

1. 正演計算
2. 梯度的反演計算
3. 將梯度應用於變量

例如，咱們能夠按如下方式 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/…