Google Colab：为什么 CPU 比 TPU 快？

我正在使用 Google colabTPU训练一个简单的Keras模型。删除分布式strategy并在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)

谢谢你的问题。

我认为这里发生的事情是一个开销问题——因为 TPU 运行在一个单独的虚拟机上（可通过grpc://$COLAB_TPU_ADDR），每次调用在 TPU 上运行模型都会产生一定量的开销，因为客户端（本例中为 Colab 笔记本）将图形发送到 TPU，然后编译并运行。与运行所需的时间相比，此开销很小。 ResNet50 适用于一个时期，但与运行示例中的简单模型相比要大一些。

为了在 TPU 上获得最佳效果，我们建议使用tf.data.数据集 https://www.tensorflow.org/api_docs/python/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 指南 https://www.tensorflow.org/guide/tpu，它为 MNIST 数据集提供了一个更大的图像分类模型。