TensorFlow 2.0高效開發指南

Effective TensorFlow 2.0

爲使TensorFLow用戶更高效，TensorFlow 2.0中進行了多出更改。TensorFlow 2.0刪除了篇冗餘API，使API更加一致（統一RNNs, 統一優化器），並經過Eager execution更好地與Python集成。

許多RFCs已經解釋了TensorFlow 2.0帶來的變化。本指南介紹了TensorFlow 2.0應該怎麼進行開發。這假設您已對TensorFlow 1.x有必定了解。

A brief summary of major changes

API Cleanup

許多API在TF 2.0中進行了移動或刪除。一些主要的變化包括刪除tf.app，tf.flags，使tf.logging支持如今開源的absl-py，從新生成項目的tf.contribe，經過清理tf.*中那些較少使用的命名空間，例如tf.math。一些API已替換爲本身的2.0版本-tf.summary,tf.keras.metrics, 和tf.keras.optimizers。最快升級應用這些重命名帶來的變化可以使用v2升級腳本。

Eager execution

TensorFlow 1.x要求用戶經過tf.*API手動的將抽象語法樹（圖）拼接在一塊兒。而後它要求用戶經過一組輸入、輸出張量傳遞給session.run()從而手動編譯調用這個圖。TensorFlow 2.0 Eager execution能夠像Python那樣執行，在2.0中，graph 和 session會像實現細節同樣。

值得注意的是tf.control_dependencies()再也不須要了，由於全部代碼都是行順序執行的（用tf.function聲明）。

No more globals

TensorFlow 1.x嚴重依賴隱式全局命名空間。當你調用tf.Variable()，它會被放入默認圖中，即便你忘了指向它的Python變量，它也會被保留在那裏。而後你能夠恢復它，但前提是你得知道它建立時的名稱。若是你沒法控制變量的建立，這很難作到。其結果是，各類各樣的機制，試圖幫助用戶再次找到他們的變量，以及爲框架找到用戶建立的變量：Variable scopes, global collections。例如tf.get_global_step()，tf.global_variables_initializer()，還有優化器隱式計算全部可訓練變量的梯度等等。
TensorFlow 2.0消除了這些機制（Variable 2.0 RFC）默認支持的機制：跟蹤你的變量！若是你忘記了一個tf.Variable，它就會看成垃圾被回收。

Functions, not sessions

session.run()幾乎能夠像函數同樣調用：指定輸入和被調用的函數，你能夠獲得一組輸出。在TensorFlow 2.0中，您可使用Python函數tf.function()來標記它以進行JIT編譯，以便TensorFlow將其做爲單個圖運行(Function 2.0 RFC)。這種機制容許TensorFlow 2.0得到圖模型全部的好處：

• 性能：函數能夠被優化（node pruning, kernel fusion, etc.）
• 可移植性：該功能能夠被導出/從新導入（SavedModel 2.0 RFC），容許用戶重用和共享模塊化TensorFlow功能。

# TensorFlow 1.X
outputs = session.run(f(placeholder), feed_dict={placeholder: input})
# TensorFlow 2.0
outputs = f(input)

憑藉穿插Python 和TensorFlow代碼的能力，咱們但願用戶可以充分利用Python的表現力。除了在沒有Python解釋器的狀況下執行TensorFlow，如mobile, C++, 和 JS。爲了幫助用戶避免在添加時重寫代碼@tf.function， AutoGraph會將Python構造的一個子集轉換爲他們的TensorFlow等價物：

• for/while -> tf.while_loop (支持break 和 continue)
• if->tf.cond
• for _ in dataset -> dataset.reduce

AutoGraph支持控制流的任意嵌套，這使得能夠有較好性能而且簡潔地實現許多複雜的ML程序，如序列模型，強化學習，自定義訓練循環等。

Recommendations for idiomatic TensorFlow 2.0

Refactor your code into smaller functions

TensorFlow 1.x中常見使用模式是「kitchen sink」策略，其中全部可能的計算的聯合被預先佈置，而後選擇被評估的張量，經過session.run()運行。在TensorFlow 2.0中，用戶應該將代碼重構爲較小的函數，這些函數根據須要被調用。一般，沒有必要用tf.function去裝飾那些比較小的函數；僅用tf.function去裝飾高等級的計算，例如，訓練的一個步驟，或模型的前向傳遞。

Use Keras layers and models to manage variables

Keras模型和圖層提供了方便variables和 trainable_variables屬性，它以遞歸方式收集全部因變量。這使得在本地管理變量很是容易。

對比：

def dense(x, W, b):
  return tf.nn.sigmoid(tf.matmul(x, W) + b)

@tf.function
def multilayer_perceptron(x, w0, b0, w1, b1, w2, b2 ...):
  x = dense(x, w0, b0)
  x = dense(x, w1, b1)
  x = dense(x, w2, b2)
  ...

# 你仍然須要管理w_i和b_i，它們的形狀遠離代碼定義。

Keras版本：

# 能夠調用每一個圖層，其簽名等效於 linear(x)
layers = [tf.keras.layers.Dense(hidden_size, activation=tf.nn.sigmoid) for _ in range(n)]
perceptron = tf.keras.Sequential(layers)

# layers[3].trainable_variables => returns [w3, b3]
# perceptron.trainable_variables => returns [w0, b0, ...]

Keras layers/models繼承自tf.train.Checkpointable並集成了@tf.function，這使得直接從Keras對象導出SavedModels或checkpoint成爲可能。您不必定要使用Keras的.fitAPI來利用這些集成。

這是一個遷移學習的例子，演示了Keras如何輕鬆收集相關變量的子集。假設你正在訓練一個帶有共享主幹的多頭模型：

trunk = tf.keras.Sequential([...])
head1 = tf.keras.Sequential([...])
head2 = tf.keras.Sequential([...])

path1 = tf.keras.Sequential([trunk, head1])
path2 = tf.keras.Sequential([trunk, head2])

# Train on primary dataset
for x, y in main_dataset:
  with tf.GradientTape() as tape:
    prediction = path1(x)
    loss = loss_fn_head1(prediction, y)
  # Simultaneously optimize trunk and head1 weights.
  gradients = tape.gradients(loss, path1.trainable_variables)
  optimizer.apply_gradients(gradients, path1.trainable_variables)

# Fine-tune second head, reusing the trunk
for x, y in small_dataset:
  with tf.GradientTape() as tape:
    prediction = path2(x)
    loss = loss_fn_head2(prediction, y)
  # Only optimize head2 weights, not trunk weights
  gradients = tape.gradients(loss, head2.trainable_variables)
  optimizer.apply_gradients(gradients, head2.trainable_variables)

# You can publish just the trunk computation for other people to reuse.
tf.saved_model.save(trunk, output_path)

Combine tf.data.Datasets and @tf.function

在內存中迭代擬合訓練數據時，能夠隨意使用常規的Python迭代。或者，tf.data.Dataset是從硬盤讀取訓練數據流的最好方法。Datasets是可迭代的（不是迭代器），它能夠像在Eager模式下的其餘Python迭代同樣工做。您能夠經過用tf.function()包裝代碼來充分利用數據集異步預取/流功能，這將使用AutoGraph等效的圖操做替換Python的迭代。

@tf.function
def train(model, dataset, optimizer):
  for x, y in dataset:
    with tf.GradientTape() as tape:
      prediction = model(x)
      loss = loss_fn(prediction, y)
    gradients = tape.gradients(loss, model.trainable_variables)
    optimizer.apply_gradients(gradients, model.trainable_variables)

若是您使用Keras.fit()API，則無需擔憂數據集迭代。

model.compile(optimizer=optimizer, loss=loss_fn)
model.fit(dataset)

Take advantage of AutoGraph with Python control flow

AutoGraph提供了一種將依賴於數據的控制流轉換爲等效圖形模式的方法，如tf.cond和tf.while_loop。

數據相關控制流出現的一個常見位置是序列模型。tf.keras.layers.RNN包裝了一個RNN cell，容許您既能夠靜態也能夠動態的循環展開。爲了演示，您能夠從新實現動態展開，以下所示：

class DynamicRNN(tf.keras.Model):

  def __init__(self, rnn_cell):
    super(DynamicRNN, self).__init__(self)
    self.cell = rnn_cell

  def call(self, input_data):
    # [batch, time, features] -> [time, batch, features]
    input_data = tf.transpose(input_data, [1, 0, 2])
    outputs = tf.TensorArray(tf.float32, input_data.shape[0])
    state = self.cell.zero_state(input_data.shape[1], dtype=tf.float32)
    for i in tf.range(input_data.shape[0]):
      output, state = self.cell(input_data[i], state)
      outputs = outputs.write(i, output)
    return tf.transpose(outputs.stack(), [1, 0, 2]), state

有關AutoGraph功能的更詳細概述，請參閱指南

Use tf.metrics to aggregate data and tf.summary to log it

要記錄摘要，請使用tf.summary.(scalar|histogram|...)上下文管理器將其重定向到編寫器。（若是省略上下文管理器，則不會發生任何事情。）與TF 1.x不一樣，摘要直接發送給編寫器; 沒有單獨的「合併」操做，也沒有單獨的add_summary()調用，這意味着step必須在調用點提供該值。

summary_writer = tf.summary.create_file_writer('/tmp/summaries')
with summary_writer.as_default():
  tf.summary.scalar('loss', 0.1, step=42)

要在將數據記錄爲摘要以前聚合數據，請使用tf.metrics。Metrics是有狀態的；它們積累值並在您調用.result()時返回結果。清除積累值，請使用.reset_states()。

def train(model, optimizer, dataset, log_freq=10):
  avg_loss = tf.keras.metrics.Mean(name='loss', dtype=tf.float32)
  for images, labels in dataset:
    loss = train_step(model, optimizer, images, labels)
    avg_loss.update_state(loss)
    if tf.equal(optimizer.iterations % log_freq, 0):
      tf.summary.scalar('loss', avg_loss.result(), step=optimizer.iterations)
      avg_loss.reset_states()

def test(model, test_x, test_y, step_num):
  loss = loss_fn(model(test_x), test_y)
  tf.summary.scalar('loss', loss, step=step_num)

train_summary_writer = tf.summary.create_file_writer('/tmp/summaries/train')
test_summary_writer = tf.summary.create_file_writer('/tmp/summaries/test')

with train_summary_writer.as_default():
  train(model, optimizer, dataset)

with test_summary_writer.as_default():
  test(model, test_x, test_y, optimizer.iterations)

經過將TensorBoard指向摘要日誌目錄來可視化生成的摘要：tensorboard --logdir /tmp/summaries。

閱讀原文

• 歡迎關注個人公衆號，一塊兒學習！