# 十九、張量處理單元 **張量處理單元**（**TPU**）是**專用集成電路**（**ASIC**），它實現了針對計算要求而優化的硬件電路深度神經網絡。 TPU 基于**復雜指令集計算機**（**CISC**）指令集，該指令集實現用于訓練深度神經網絡的復雜任務的高級指令。 TPU 架構的核心在于優化矩陣運算的脈動數組。 The Architecture of TPUImage from:?https://cloud.google.com/blog/big-data/2017/05/images/149454602921110/tpu-15.png TensorFlow 提供了一個編譯器和軟件棧，可將 API 調用從 TensorFlow 圖轉換為 TPU 指令。以下框圖描述了在 TPU 棧頂部運行的 TensorFlow 模型的架構：  圖像來自[這個鏈接](https://cloud.google.com/blog/big-data/2017/05/images/149454602921110/tpu-2.png)。TPU 架構的更多信息，閱讀[博客](https://cloud.google.com/blog/big-data/2017/05/an-in-depth-look-at-googles-first-tensor-processing-unit-tpu)。 TPU 的 TensorFlow API 位于`tf.contrib.tpu`模塊中。為了在 TPU 上構建模型，使用以下三個 TPU 特定的 TensorFlow 模塊： * `tpu_config`：`tpu_config`模塊允許您創建配置對象，其中包含有關將運行模型的主機的信息。 * `tpu_estimator`：`tpu_estimator`模塊將估計器封裝在`TPUEstimatorSpec`類中。要在 TPU 上運行估計器，我們創建此類的對象。 * `tpu_optimizer`：`tpu_optimizer`模塊包裝優化器。例如，在下面的示例代碼中，我們將`tpu_optimizer`類中的 SGD 優化器包裝在`tpu_optimizer`類中。 例如，以下代碼使用 TFEstimator API 為 TPU 上的 MNIST 數據集構建 CNN 模型： 以下代碼改編自[這個頁面](https://github.com/tensorflow/tpu-demos/blob/master/cloud_tpu/models/mnist/mnist.py)。 ```py import tensorflow as tf from tensorflow.contrib.tpu.python.tpu import tpu_config from tensorflow.contrib.tpu.python.tpu import tpu_estimator from tensorflow.contrib.tpu.python.tpu import tpu_optimizer learning_rate = 0.01 batch_size = 128 def metric_fn(labels, logits): predictions = tf.argmax(logits, 1) return { "accuracy": tf.metrics.precision( labels=labels, predictions=predictions), } def model_fn(features, labels, mode): if mode == tf.estimator.ModeKeys.PREDICT: raise RuntimeError("mode {} is not supported yet".format(mode)) input_layer = tf.reshape(features, [-1, 28, 28, 1]) conv1 = tf.layers.conv2d( inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same", activation=tf.nn.relu) pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2) conv2 = tf.layers.conv2d( inputs=pool1, filters=64, kernel_size=[5, 5], padding="same", activation=tf.nn.relu) pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2) pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64]) dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu) dropout = tf.layers.dropout( inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN) logits = tf.layers.dense(inputs=dropout, units=10) onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10) loss = tf.losses.softmax_cross_entropy( onehot_labels=onehot_labels, logits=logits) if mode == tf.estimator.ModeKeys.EVAL: return tpu_estimator.TPUEstimatorSpec( mode=mode, loss=loss, eval_metrics=(metric_fn, [labels, logits])) # Train. decaying_learning_rate = tf.train.exponential_decay(learning_rate, tf.train.get_global_step(), 100000,0.96) optimizer = tpu_optimizer.CrossShardOptimizer( tf.train.GradientDescentOptimizer( learning_rate=decaying_learning_rate)) train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step()) return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op) def get_input_fn(filename): def input_fn(params): batch_size = params["batch_size"] def parser(serialized_example): features = tf.parse_single_example( serialized_example, features={ "image_raw": tf.FixedLenFeature([], tf.string), "label": tf.FixedLenFeature([], tf.int64), }) image = tf.decode_raw(features["image_raw"], tf.uint8) image.set_shape([28 * 28]) image = tf.cast(image, tf.float32) * (1\. / 255) - 0.5 label = tf.cast(features["label"], tf.int32) return image, label dataset = tf.data.TFRecordDataset( filename, buffer_size=FLAGS.dataset_reader_buffer_size) dataset = dataset.map(parser).cache().repeat() dataset = dataset.apply( tf.contrib.data.batch_and_drop_remainder(batch_size)) images, labels = dataset.make_one_shot_iterator().get_next() return images, labels return input_fn # TPU config master = 'local' #URL of the TPU instance model_dir = '/home/armando/models/mnist' n_iterations = 50 # number of iterations per TPU training loop n_shards = 8 # number of TPU chips run_config = tpu_config.RunConfig( master=master, evaluation_master=master, model_dir=model_dir, session_config=tf.ConfigProto( allow_soft_placement=True, log_device_placement=True ), tpu_config=tpu_config.TPUConfig(n_iterations, n_shards ) ) estimator = tpu_estimator.TPUEstimator( model_fn=model_fn, use_tpu=True, train_batch_size=batch_size, eval_batch_size=batch_size, config=run_config) train_file = '/home/armando/datasets/mnist/train' # input data file train_steps = 1000 # number of steps to train for estimator.train(input_fn=get_input_fn(train_file), max_steps=train_steps ) eval_file = '/home/armando/datasets/mnist/test' # test data file eval_steps = 10 estimator.evaluate(input_fn=get_input_fn(eval_file), steps=eval_steps ) ``` 有關在 TPU 上構建模型的更多示例，[請訪問此鏈接](https://github.com/tensorflow/tpu-demos)。