如何為 Cloud TPU 編寫自定義估算器模型
文 / Google Cloud Platform 技術主管 Lak Lakshmanan (@lak_gcp)
來源 | TensorFlow 公眾號
張量處理單元 (TPU) 可加速處理 Google 內各種機器學習工作負載，並可供 Google Cloud 客戶使用。您可以在 Cloud TPU 參考模型儲存區找到啟用 TPU 的頂尖影象模型版本，例如 ResNet 和 AmoebaNet；您還可以使用強大的 Tensor2Tensor 庫，在 TPU 上執行文字摘要和問答任務。這些教程會為您分步介紹如何使用當下很多最熱門的 Cloud TPU 參考模型。
注：儲存區連結 ofollow,noindex">github.com/tensorflow/…
教程連結 cloud.google.com/tpu/docs/tu…
但如果您擁有自定義 TensorFlow 模型，又該如何做呢？在本文中，我會逐步介紹編寫自定義估算器以便在 Cloud TPU 上執行的全過程。在此過程中，我會指出需要注意的地方和建議採用的最佳實踐。您可以在 GitHub 上找到此解決方案的完整程式碼；本文僅列出相關程式碼片段。
注：解決方案的完整程式碼連結 github.com/GoogleCloud…
自定義 TensorFlow 估算器包含以模型函式傳遞的基類估算器：
1def train_and_evaluate(output_dir, nsteps): 2estimator = tf.estimator.Estimator( 3model_fn = model_fn, 4model_dir = output_dir) 複製程式碼
模型函式會接收特徵、標籤和模式，並返回 EstimatorSpec。例如，影象分類問題的模型函式可能包含
1def model_fn(features, labels, mode):
2# write the model to compute predictions, loss, etc. from the model
3
4return tf.estimator.EstimatorSpec(
5mode=mode,
6predictions={"probabilities": probabilities,
7"classid": class_int, "class": class_str},
8loss=loss,
9train_op=train_op,
10eval_metric_ops=evalmetrics,
11export_outputs={'classes': tf.estimator.export.PredictOutput(
12{"probabilities": probabilities, "classid": class_int,
13"class": class_str})}
14)
複製程式碼
TensorFlow 中的 tf.contrib.tpu 包提供了包裝器類，可助您以適當方式編寫程式碼，以便在 CPU、GPU 和 Cloud TPU 上執行程式碼。下面我們就來看看如何以這種不受加速器限制的方式編寫自定義估計器。
1.將輸入資料轉換為 TF 記錄
Cloud TPU 的速度非常快，一不小心，您的訓練就會變成以讀寫（“饋入” 和 “饋出”）資料和儲存檢查點為主。讓 TPU 等待輸入/輸出會造成浪費，所以我們會做幾件事以充分利用 TPU 用於計算的時間。
首先是避免在估算器的輸入函式中進行資料解析和整理，而是預先將資料轉換為 TF 記錄。與單個影象檔案相比，批量處理 TF 記錄更為簡單，因為記錄本身包含標籤，如此可以減少系統必須讀取的小型檔案數量。我使用 Apache Beam 進行這種轉換。您可以在官方 TPU 儲存區找到讀取 JPEG 和編寫 TF 記錄的指令碼。您可以在 Cloud Dataflow 上大規模地執行 Apache Beam 程式，但如果您的資料來源目前不在 Google Cloud 上，則只能在大型 VM 上本地執行此程式（請務必用 pip 安裝 apache-beam）。
注：JPEG 和編寫 TF 記錄連結 github.com/tensorflow/…
TF 記錄是字典。對於影象分類，上述管道編寫的兩個條目很重要，分別是：“image/class/label”（採用 int64）和 “image/encoded”（由 JPEG 檔案內容組成）。
2.編寫輸入函式以讀取 TF 記錄
與任何估算器一樣，您需要編寫輸入函式，以讀取這些 TF 記錄。使用 Dataset API 可極大地簡化此任務，但還需注意幾個問題。在講解過程中，我會指出這些問題。
以下是我的輸入函式：
1def make_input_fn(pattern, mode, num_cores=8, transpose_input=False): 2def _set_shapes(batch_size, images, labels): 3"""Statically set the batch_size dimension.""" 4if transpose_input: 5images.set_shape(images.get_shape().merge_with( 6tf.TensorShape([None, None, None, batch_size]))) 7labels.set_shape(labels.get_shape().merge_with( 6tf.TensorShape([batch_size]))) 9else: 10images.set_shape(images.get_shape().merge_with( 11tf.TensorShape([batch_size, None, None, None]))) 12labels.set_shape(labels.get_shape().merge_with( 13tf.TensorShape([batch_size]))) 14return images, labels 15 16def _input_fn(params): 17batch_size = params['batch_size'] 18is_training = (mode == tf.estimator.ModeKeys.TRAIN) 19 20# read the dataset 21dataset = tf.data.Dataset.list_files(pattern, shuffle=is_training) 22if is_training: 23dataset = dataset.repeat() 24def fetch_dataset(filename): 25buffer_size = 8 * 1024 * 1024 # 8 MiB per file 26dataset = tf.data.TFRecordDataset(filename, buffer_size=buffer_size) 27return dataset 28dataset = dataset.apply( 29tf.contrib.data.parallel_interleave( 30fetch_dataset, cycle_length=64, sloppy=True)) 31dataset = dataset.shuffle(1024) 32 33# augment and batch 34dataset = dataset.apply( 35tf.contrib.data.map_and_batch( 36read_and_preprocess, batch_size=batch_size, 37num_parallel_batches=num_cores, drop_remainder=True 38)) 39 40if transpose_input: 41dataset = dataset.map( 42lambda images, labels: (tf.transpose(images, [1, 2, 3, 0]), labels), 43num_parallel_calls=num_cores) 44 45# assign static shape 46dataset = dataset.map( 47functools.partial(_set_shapes, batch_size) 48) 49 50# prefetch data while training 51dataset = dataset.prefetch(tf.contrib.data.AUTOTUNE) 52return dataset 53 54return _input_fn 複製程式碼
請注意，輸入函式採用 params 引數。實際上，這將是傳遞至訓練程式的命令列引數，如此一來，我們便可提取有關資料集的詳情，例如訓練次數和評估影象。
batch_size 很特別，因為 TPU 有多個核心，而 batch_size 由 TPU 估算器設定，且為有效批次大小。您必須完全返回 batch_size 記錄，而不能返回部分填充的批次。由於您會無限期地迴圈使用訓練資料，所以在訓練期間不會出現此問題。但這意味著最簡單的做法是將評估資料集向下舍入為核心數的倍數。如果核心數為 8，而評估資料集中有 1026 張影象，您只能使用前 1024 張影象進行評估。剩餘的 2 張影象則會捨棄。（我們也有方法在 Cloud TPU 中處理最後剩下的部分批次，我就不在此贅述。）
與任何分散式訓練一樣，您應確保每個工作器看到不同的資料子集，這是由所有檔案的並行交錯及緩衝區本身內部的記錄重排進行處理。
影象分類問題的一個常見需求是通過新增隨機裁剪及翻轉等方法來增強原始資料。我通過 read_and_preprocess 函式做到這一點。請注意，我將此函式應用於各個 TF 記錄並建立了 8 個並行批次，同時捨棄剩餘的任何記錄（再次提醒，這在訓練期間不會造成任何影響，因為您會無限期重複進行訓練）。
接下來是轉置資料。事實證明，在 TPU 中轉置資料以保持批次大小可以極大地提高效能。因此，我們可以根據需要採取此做法。如果我們在 GPU 或 CPU 上執行程式，則 transpose_input 標記會變為 false。
TPU 需要靜態大小的張量。儘管我們已確保維持這種情況（通過捨棄剩餘的批次），但仍需為核心 TensorFlow 編寫 Dataset API，這是更常見的做法。因此，我們呼叫一個函式，將資料集中的 batch_size 從 None 改為 batch_size。
最後的優化操作至關重要。我們需要預取資料。換句話說，當 TPU 處理一批記錄時，我們會通過 I/O 執行緒尋找並提取下一批次。如此一來，我們便可最大限度地利用 TPU（或 GPU），而且這對 CPU 沒有任何影響。
3.處理 TF 記錄
（上述）輸入函式會設定處理輸入內容的方式，但實際的解析操作還需由名為 read_and_preprocess() 的方法執行。此方法如下所示：
1def read_and_preprocess(example_data):
2parsed = tf.parse_single_example(example_data, {
3'image/encoded': tf.FixedLenFeature((), tf.string, ''),
4'image/class/label': tf.FixedLenFeature([], tf.int64, 1),
5})
6image_bytes = tf.reshape(parsed['image/encoded'], shape=[])
7label = tf.cast(
8tf.reshape(parsed['image/class/label'], shape=[]), dtype=tf.int32) - 1
9
10# end up with pixel values that are in the -1, 1 range
11image = tf.image.decode_jpeg(image_bytes, channels=NUM_CHANNELS)
12image = tf.image.convert_image_dtype(image, dtype=tf.float32) # 0-1
13image = tf.expand_dims(image, 0) # resize_bilinear needs batches
14
15image = tf.image.resize_bilinear(
16image, [HEIGHT + 10, WIDTH + 10], align_corners=False)
17image = tf.squeeze(image)# remove batch dimension
18image = tf.random_crop(image, [HEIGHT, WIDTH, NUM_CHANNELS])
19image = tf.image.random_flip_left_right(image)
20image = tf.image.random_brightness(image, max_delta=63.0 / 255.0)
21image = tf.image.random_contrast(image, lower=0.2, upper=1.8)
22
23
24#pixel values are in range [0,1], convert to [-1,1]
25image = tf.subtract(image, 0.5)
26image = tf.multiply(image, 2.0)
27return image, label
複製程式碼
這裡有兩個重要的注意事項。第一是使用 parse_single_example，因為我們是從 map() 呼叫此函式，所以會針對單個 TF 記錄呼叫。我們從記錄中提取相關資訊（經過編碼的影象和標籤），然後將其用於構建必要的張量。第二個注意事項是，這些資料必須為數值。比如，我無法傳回標籤字串，因為 TPU 只能處理數值型資料。我們需要在預處理管道中計算標籤索引，標籤此時只會是整數。
4.提供輸入函式
訓練模型之後，您需要部署此模型，並通過 TF Serving 提供。以下程式碼與您使用任何其他估算器時要用到的程式碼相同：
1def serving_input_fn():
2# Note: only handles one image at a time
3feature_placeholders = {'image_bytes':
4tf.placeholder(tf.string, shape=())}
5image, _ = read_and_preprocess(
6tf.squeeze(feature_placeholders['image_bytes']))
7features = {
8'image': tf.expand_dims(image, 0)
9}
10return tf.estimator.export.ServingInputReceiver(features, feature_placeholders)
複製程式碼
TPU 已針對批次推理進行優化；如果您的用例需要線上預測，目前最好是通過 CPU 或 GPU 提供（根據批次大小和模型複雜程度而定）。編寫輸入函式時，我假定自己只傳送一張影象，所以實際上是指通過 CPU/GPU 提供。
5.模型函式
模型函式需要建立並返回 TPUEstimatorSpec。實現方式如下所示：
1def image_classifier(features, labels, mode, params):
2image = features
3if isinstance(features, dict):
4image = features['image']
5
6ylogits, nclasses = cnn_model(image, mode, params)
7
8probabilities = tf.nn.softmax(ylogits)
9class_int = tf.cast(tf.argmax(probabilities, 1), tf.int32)
10class_str = tf.gather(LIST_OF_LABELS, class_int)
11
12if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
13loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(
14logits=ylogits, labels=tf.one_hot(labels, nclasses)))
15
16def metric_fn(class_int, labels):
17return {'accuracy': tf.metrics.accuracy(class_int, labels)}
18evalmetrics = (metric_fn, [class_int, labels])
19
20if mode == tf.estimator.ModeKeys.TRAIN:
21# this is needed for batch normalization, but has no effect otherwise
22update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
23optimizer = tf.train.AdamOptimizer(learning_rate=params['learning_rate'])
24if params['use_tpu']:
25optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer) # TPU change 1
26with tf.control_dependencies(update_ops):
27train_op = optimizer.minimize(loss, tf.train.get_global_step())
28else:
29train_op = None
30else:
31loss = None
32train_op = None
33evalmetrics = None
34
35return tf.contrib.tpu.TPUEstimatorSpec(# TPU change 2
36mode=mode,
37predictions={"probabilities": probabilities,
38"classid": class_int, "class": class_str},
39loss=loss,
40train_op=train_op,
41eval_metrics=evalmetrics,
42export_outputs={'classes': tf.estimator.export.PredictOutput(
43{"probabilities": probabilities, "classid": class_int,
44"class": class_str})}
45)
複製程式碼
所傳入的特徵可能會是影象（我的訓練和評估輸入函式）或字典（我的提供輸入函式）。我進行檢查並從特徵中檢索影象。
然後，我在影象上呼叫實際模型數學函式。這應該是使用 tf.layers 的常見 TensorFlow 程式碼。請瀏覽完整的原始碼以瞭解其形式。
由於這是分類問題，所以我使用 softmax 並根據各個類別的 logit 計算整數標籤和字串標籤，之後使用了 argmax 和 gather。我還計算了交叉熵損失，這與任何其他估算器類似。
其中一個區別在於，一般的估算器需要將評估指標用作字典，而 TPUEstimator 需要能夠在控制 CPU 或 TPU 上呼叫的函式。因此，指定 eval 指標的方式稍有不同。
如果您使用 TPU，則所用的優化器必須包裝在 CrossShardOptimizer 中。這樣可以在不同核心中分配優化任務。
訓練操作就是將此交叉碎片優化損失最小化。請使用 optimizer.minimize()，而非 layers.optimize_loss()。
將上述所有操作整合在一起後，系統會返回 TPU 估算器規範。
6.訓練和評估迴圈
您可能很熟悉估算器的 train_and_evaluate 迴圈。可惜此迴圈（尚）無法與 TPU 有效配合使用。幸運的是，建立您自己的迴圈並不太難，這讓您可以更好地控制檢查頻率和內容（回想這樣的情景，您想盡可能減少過於頻繁的檢查導致的環境切換和 I/O 開銷）。
1def train_and_evaluate(output_dir, hparams):
2STEPS_PER_EVAL = 1000
3max_steps = hparams['train_steps']
4eval_batch_size = min(1024, hparams['num_eval_images'])
5eval_batch_size = eval_batch_size - eval_batch_size % 8# divisible by num_cores
6tf.logging.info('train_batch_size=%deval_batch_size=%dmax_steps=%d',
7hparams['train_batch_size'],
8eval_batch_size,
9max_steps)
10
11# TPU change 3
12if hparams['use_tpu']:
13tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
14hparams['tpu'],
15zone=hparams['tpu_zone'],
16project=hparams['project'])
17config = tf.contrib.tpu.RunConfig(
18cluster=tpu_cluster_resolver,
19model_dir=output_dir,
20save_checkpoints_steps=STEPS_PER_EVAL,
21tpu_config=tf.contrib.tpu.TPUConfig(
22iterations_per_loop=STEPS_PER_EVAL,
23per_host_input_for_training=True))
24else:
25config = tf.contrib.tpu.RunConfig()
26
27estimator = tf.contrib.tpu.TPUEstimator(# TPU change 4
28model_fn=image_classifier,
29config=config,
30params=hparams,
31model_dir=output_dir,
32train_batch_size=hparams['train_batch_size'],
33eval_batch_size=eval_batch_size,
34use_tpu=hparams['use_tpu']
35)
複製程式碼
首先，提取一些命令列引數，並用其指定最大步數以及訓練和評估的批次大小。
接下來是尋找 TPU。如果您已在 Google 計算引擎上自行建立了 Cloud TPU，可能已為其命名。假設此名稱（“tpu”）作為命令列引數傳入。如果您使用 Cloud ML Engine，系統會自動推斷 TPU 名稱和區域等內容。請務必僅在已設定 use_tpu 標記的情況下執行此操作。如果使用者是在 CPU 或 GPU 上執行程式，則只需建立空白 RunConfig。
接下來，使用模型函式、配置、引數和批次大小建立 TPUEstimator。建立估算器後，我們便可進入真實訓練和評估迴圈：
1# load last checkpoint and start from there
2current_step = load_global_step_from_checkpoint_dir(output_dir)
3steps_per_epoch = hparams['num_train_images'] // hparams['train_batch_size']
4tf.logging.info('Training for %d steps (%.2f epochs in total). Current'
5' step %d.',
6max_steps,
7max_steps / steps_per_epoch,
8current_step)
9
10start_timestamp = time.time()# This time will include compilation time
11
12while current_step < hparams['train_steps']:
13# Train for up to steps_per_eval number of steps.
14# At the end of training, a checkpoint will be written to --model_dir.
15next_checkpoint = min(current_step + STEPS_PER_EVAL, max_steps)
16estimator.train(input_fn=train_input_fn, max_steps=next_checkpoint)
17current_step = next_checkpoint
18tf.logging.info('Finished training up to step %d. Elapsed seconds %d.',
19next_checkpoint, int(time.time() - start_timestamp))
20
21# Evaluate the model on the most recent model in --model_dir.
22# Since evaluation happens in batches of --eval_batch_size, some images
23# may be excluded modulo the batch size. As long as the batch size is
24# consistent, the evaluated images are also consistent.
25tf.logging.info('Starting to evaluate at step %d', next_checkpoint)
26eval_results = estimator.evaluate(
27input_fn=eval_input_fn,
28steps=hparams['num_eval_images'] // eval_batch_size)
29tf.logging.info('Eval results at step %d: %s', next_checkpoint, eval_results)
30
31elapsed_time = int(time.time() - start_timestamp)
32tf.logging.info('Finished training up to step %d. Elapsed seconds %d.',
33max_steps, elapsed_time)
複製程式碼
TensorFlow 估算器的運作方式是從先前已有的檢查點執行暖啟動。我們可以載入輸出目錄中提供的檢查點，以進行復制。然後，我們會一次性逐步完成訓練資料 train_batch_size 步驟，直至達到所指定的最大步數。
在本文的例子中，我對完整評估資料集中的每個檢查點都進行了評估，但顯然，您可以減少此訓練的計算量。
7.匯出模型以供使用
最後，在完成訓練後，我匯出已儲存的模型。您可以使用 TF Serving 或 Cloud ML Engine 來部署已儲存的模型，以進行預測。
1# export similar to Cloud ML Engine / TF Serving convention
2tf.logging.info('Starting to export model.')
3estimator.export_savedmodel(
4export_dir_base=os.path.join(output_dir, 'export/exporter'),
5serving_input_receiver_fn=serving_input_fn)
複製程式碼
此時，我們便有了一個可以在 Cloud TPU 上訓練的自定義估算器模型。採用這種方式編寫模型（例如使用 use_tpu 標記並提供轉置為可選項），同樣的程式碼也支援各種不同的硬體，包括 CPU 和 GPU 在內。因此，我們的估算器模型實際可用於全部的三類硬體。
後續步驟：
從 GitHub 下載本文隨附的程式碼，然後進行試用 執行程式碼實驗室，瞭解如何在 TPU 上執行 ResNet 訓練您自己的資料（無需編寫任何程式碼）
注：程式碼連結 github.com/GoogleCloud…
在 TPU 上執行 ResNet 連結 codelabs.developers.google.com/codelabs/tp…
在 Coursera 上參加使用 TensorFlow 進行機器學習專業課程；此課程會逐步介紹 TensorFlow 概念，以及如何在 Google Cloud 上大規模地訓練、調整和部署 ML 模型。
注：使用 TensorFlow 進行機器學習連結 www.coursera.org/specializat…