Google Colab：为什么 CPU 比 TPU 快？

Question

我正在使用 Google colab TPU来训练一个简单的Keras模型。去掉分布式策略，在CPU上运行相同的程序比TPU快得多。这怎么可能？

import timeit
import os
import tensorflow as tf
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam

# Load Iris dataset
x = load_iris().data
y = load_iris().target

# Split data to train and validation set
x_train, x_val, y_train, y_val = train_test_split(x, y, test_size=0.30, shuffle=False)

# Convert train data type to use TPU 
x_train = x_train.astype('float32')
x_val = x_val.astype('float32')

# Specify a distributed strategy to use TPU
resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.contrib.distribute.initialize_tpu_system(resolver)
strategy = tf.contrib.distribute.TPUStrategy(resolver)

# Use the strategy to create and compile a Keras model
with strategy.scope():
  model = Sequential()
  model.add(Dense(32, input_shape=(4,), activation=tf.nn.relu, name="relu"))
  model.add(Dense(3, activation=tf.nn.softmax, name="softmax"))
  model.compile(optimizer=Adam(learning_rate=0.1), loss='logcosh')

start = timeit.default_timer()

# Fit the Keras model on the dataset
model.fit(x_train, y_train, batch_size=20, epochs=20, validation_data=[x_val, y_val], verbose=0, steps_per_epoch=2)

print('\nTime: ', timeit.default_timer() - start)

Answer 1

谢谢你的问题。

我认为，这里发生了什么是架空的事-因为在一个单独的虚拟机（在访问的TPU运行grpc://$COLAB_TPU_ADDR），每次调用上运行TPU模型招致开销的一定量的客户端（在这种情况下Colab笔记本）将图发送到 TPU，然后编译并运行。与运行 ResNet50 等一个时期所需的时间相比，这种开销很小，但与运行像您示例中的简单模型这样的简单模型相比却很大。

为了在 TPU 上获得最佳结果，我们建议使用tf.data.Dataset。我为 TensorFlow 2.2 更新了您的示例：

%tensorflow_version 2.x
import timeit
import os
import tensorflow as tf
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam

# Load Iris dataset
x = load_iris().data
y = load_iris().target

# Split data to train and validation set
x_train, x_val, y_train, y_val = train_test_split(x, y, test_size=0.30, shuffle=False)

# Convert train data type to use TPU 
x_train = x_train.astype('float32')
x_val = x_val.astype('float32')

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(20)
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(20)

# Use the strategy to create and compile a Keras model
with strategy.scope():
  model = Sequential()
  model.add(Dense(32, input_shape=(4,), activation=tf.nn.relu, name="relu"))
  model.add(Dense(3, activation=tf.nn.softmax, name="softmax"))
  model.compile(optimizer=Adam(learning_rate=0.1), loss='logcosh')

start = timeit.default_timer()

# Fit the Keras model on the dataset
model.fit(train_dataset, epochs=20, validation_data=val_dataset)

print('\nTime: ', timeit.default_timer() - start)

这大约需要 30 秒才能运行，而在 CPU 上运行大约需要 1.3 秒。我们可以通过重复数据集并运行一个长时期而不是几个小时期来大大减少这里的开销。我用这个替换了数据集设置：

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).repeat(20).batch(20)
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(20)

并fit用这个替换了电话：

model.fit(train_dataset, validation_data=val_dataset)

这使我的运行时间减少到大约 6 秒。这仍然比 CPU 慢，但对于这样一个可以轻松在本地运行的小模型来说，这并不奇怪。一般来说，您会发现将 TPU 用于更大的模型会带来更多好处。我建议查看TensorFlow 的官方 TPU 指南，它为 MNIST 数据集提供了一个更大的图像分类模型。

Answer 2

这可能是由于您使用的批量大小所致。与 CPU 和 GPU 相比，TPU 的训练速度高度依赖于批量大小。请访问以下网站了解更多信息： https://cloud.google.com/tpu/docs/performance-guide

Cloud TPU 硬件不同于 CPU 和 GPU。从较高的层面来看，CPU 的特点是具有少量的高性能线程。GPU 的特点是具有大量低性能线程。具有 128 x 128 矩阵单元的 Cloud TPU 可以被视为一个非常强大的线程（每个周期可以执行 16K 操作），也可以被视为以管道方式连接的 128 x 128 个微小的简单线程。相应地，当寻址存储器时，需要8的倍数（浮点数），对于针对矩阵单元的操作也需要128的倍数。

这意味着批量大小应该是 128 的倍数，具体取决于 TPU 的数量。Google Colab 为您提供了 8 个 TPU，因此在最好的情况下您应该选择 128 * 8 = 1024 的批量大小。