Google最近在基于
Colaboratory的基于云的机器学习
平台上免费提供了其
张量处理单元（TPU）的访问权限。 张量处理器是由Google开发的专用集成电路（ASIC），用于使用TensorFlow库进行机器学习。 我决定尝试在Keras上学习TPU卷积网络，该网络可以识别CIFAR-10图像中的物体。 完整的解决方案代码可以在
笔记本电脑上查看和运行。
图片cloud.google.com张量处理器
在Habré上已经写了TPU的布置方式（
这里 ，
这里和
这里 ），以及
为什么TPU非常适合训练神经网络 。 因此，我不会深入研究TPU架构的细节，而只会考虑训练神经网络时需要考虑的功能。
现在有三代张量处理器，最后的第三代TPU的性能为420 TFlops（每秒数万亿个浮点运算），其中包含128 GB的高带宽内存。 但是，只有第二代TPU在Colaboratory上可用，它们具有180 TFlops的性能和64 GB的内存。 将来，我将考虑这些TPU。
张量处理器由四个芯片组成，每个芯片包含两个内核，TPU中总共有八个内核。 使用复制在所有内核上并行进行TPU训练：每个内核运行TensorFlow图的一个副本，该副本的数据量为八分之一。
张量处理器的基础是矩阵单元（MXU）。 它使用128x128
脉动阵列的狡猾数据结构来高效执行矩阵运算。 因此，为了最大程度地使用TPU设备资源，小型样本或特征的维数必须是128（
源 ）的倍数。 而且，由于TPU存储系统的性质，希望微型样本和特征的尺寸为8的倍数。
协作平台
Colaboratory是Google的先进机器学习技术的云平台。 您可以免费获得带有已安装的流行库TensorFlow，Keras，sklearn，pandas等的虚拟机。 最方便的是，您可以在Colaboratory上运行类似于Jupyter的笔记本电脑。 笔记本电脑存储在Google云端硬盘中，您可以分发笔记本甚至组织协作。 这是便携式计算机在协作室中的外观（
图片可点击 ）：
您可以在笔记本电脑上的浏览器中编写代码，该代码可以在Google Cloud的虚拟机上运行。 汽车会签发给您12个小时，然后停下来。 但是，没有什么可以阻止您启动另一个虚拟机并再工作12个小时。 只需记住，虚拟机停止后，虚拟机中的所有数据都会被删除。 因此，请不要忘记将必要的数据保存到您的计算机或Google云端硬盘中，然后在重新启动虚拟机之后再次下载。
有关在Colaboratory平台上工作的详细说明，请参见
此处 ，
此处和
此处 。
将张量处理器连接到协作实验室
默认情况下，Colaboratory不使用GPU或TPU计算加速器。 您可以在菜单运行系统->更改运行系统类型->硬件加速器中连接它们。 在出现的列表中，选择“ TPU”：
选择加速器类型后，将与Colaboratory便携式计算机连接的虚拟机将重新启动，并且TPU将可用。
如果您将任何数据下载到虚拟机，则在重新启动过程中将删除该数据。 您必须再次下载数据。
用于CIFAR-10识别的Keras神经网络
举例来说，让我们尝试在TPU上训练Keras神经网络，以识别来自
CIFAR-10数据集的图像。 这是一个受欢迎的数据集，其中包含10类物体的小型图像：飞机，汽车，鸟，猫，鹿，狗，青蛙，马，船和卡车。 类不相交，图片中的每个对象仅属于一个类。
使用Keras下载CIFAR-10数据集：
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
为了创建神经网络，我得到了一个单独的函数。 我们将两次创建相同的模型：用于TPU的模型的第一个版本，我们将在该模型上进行训练，而用于CPU的第二个模型，我们将在其中识别对象。
def create_model(): input_layer = Input(shape=(32, 32, 3), dtype=tf.float32, name='Input') x = BatchNormalization()(input_layer) x = Conv2D(32, (3, 3), padding='same', activation='relu')(x) x = Conv2D(32, (3, 3), activation='relu', padding='same')(x) x = MaxPooling2D(pool_size=(2, 2))(x) x = Dropout(0.25)(x) x = BatchNormalization()(x) x = Conv2D(64, (3, 3), padding='same', activation='relu')(x) x = Conv2D(64, (3, 3), activation='relu')(x) x = MaxPooling2D(pool_size=(2, 2))(x) x = Dropout(0.25)(x) x = Flatten()(x) x = Dense(512, activation='relu')(x) x = Dropout(0.5)(x) output_layer = Dense(10, activation='softmax')(x) model = Model(inputs=[input_layer], outputs=[output_layer]) model.compile( optimizer=tf.train.AdamOptimizer(0.001), loss=tf.keras.losses.sparse_categorical_crossentropy, metrics=['sparse_categorical_accuracy']) return model
到目前为止，Keras优化器无法在TPU上使用，因此，在编译模型时，将指定TensorFlow的优化器。
我们为CPU创建Keras模型，在下一步中，我们将其转换为TPU模型：
cpu_model = create_model()
将Keras神经网络转换为TPU模型
Keras和TensorFlow上的模型可以在GPU上进行训练而无需任何更改。 到目前为止，您还不能在TPU上执行此操作，因此您必须将我们创建的模型转换为TPU的模型。
首先，您需要找出可供我们使用的TPU的位置。 在协作平台上，可以使用以下命令完成此操作：
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
在我的情况下，TPU地址原来是这样的
grpc://10.102.233.146:8470 。 对于不同的发布会，地址是不同的。
现在，您可以使用
keras_to_tpu_model函数获取TPU的模型：
tf.logging.set_verbosity(tf.logging.INFO) tpu_model = tf.contrib.tpu.keras_to_tpu_model( cpu_model, strategy=tf.contrib.tpu.TPUDistributionStrategy( tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
第一行包括“信息”级别的日志记录。 这是模型转换日志：
INFO:tensorflow:Querying Tensorflow master (b'grpc://10.102.233.146:8470') for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
...
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.您可以看到在我们之前指定的地址找到了TPU，它具有8个核心。 我们还会看到警告，
tpu_model是实验性的，可能随时更改或删除。 我希望随着时间的流逝，将有可能直接在TPU上训练Keras模型而无需进行任何转换。
我们在TPU上训练模型
可以通过调用
fit方法以通常的方式为Keras训练TPU模型：
history = tpu_model.fit(x_train, y_train, batch_size=128*8, epochs=50, verbose=2)
这里有什么功能。 我们记得为了有效地使用TPU，最小样本大小必须是128的倍数。此外，使用最小样本中所有数据的八分之一对每个TPU内核进行训练。 因此，我们将训练期间的迷你样本大小设置为128 * 8，我们为每个TPU核心获得128张图片。 您可以使用较大的大小，例如256或512，则性能会更高。
就我而言，一个时代的训练平均需要6 s。
50年代的教育质量：
Epoch 50/50
- 6s - loss: 0.2727 - sparse_categorical_accuracy: 0.9006对训练数据正确回答的比例为90.06％。 我们使用TPU检查测试数据的质量：
scores = tpu_model.evaluate(x_test, y_test, verbose=0, batch_size=batch_size * 8) print(" : %.2f%%" % (scores[1]*100))
: 80.79%现在保存经过训练的模型的权重：
tpu_model.save_weights("cifar10_model.h5")
TensorFlow将给我们一条消息，即权重已从TPU传输到CPU：
INFO:tensorflow:Copying TPU weights to the CPU应当注意，训练有素的网络的权重已保存在协作虚拟机的磁盘上。 当虚拟机停止时，来自虚拟机的所有数据将被删除。 如果您不想失去训练有素的体重，请将其保存到计算机中：
from google.colab import files files.download("cifar10_model.h5")
识别CPU上的对象
现在，让我们尝试使用在TPU上训练的模型，以便使用CPU识别图像中的对象。 为此，请再次创建模型，并将在TPU上训练的权重加载到其中：
model = create_model() model.load_weights("cifar10_model.h5")
该模型已准备就绪，可以在中央处理器上使用。 让我们尝试借助其帮助来识别CIFAR-10测试套件的图像之一：
index=111 plt.imshow(toimage(x_test[index])) plt.show()
图片很小，但是您可以理解这是一架飞机。 我们开始认识：
我们获得了神经元输出值的列表，除第一个值（对应于平面）外，几乎所有输出值都接近于零。
[[9.81738389e-01 2.91262069e-07 1.82225723e-02 9.78524668e-07
5.89265142e-07 6.76223244e-10 1.03252004e-10 9.23009047e-09
3.71878523e-05 3.16599618e-08]]
识别成功！
总结
可以在协作平台上演示TPU的可操作性，它可以用于在Keras上训练神经网络。 但是，CIFAR-10数据集太小；不足以完全加载TPU资源。 与GPU相比，加速度很小（您可以通过选择GPU作为加速器而不是TPU并再次训练模型来检查自己）。
在Habré上有一篇文章，其中介绍
了在训练网络ResNet-50时TPU和GPU V100的性能 。 在执行此任务时，TPU表现出与四个V100 GPU相同的性能。 很好，Google免费提供了如此强大的神经网络学习加速器！
演示Keras神经网络在TPU上训练的视频。
有用的链接
- 具有完整Keras TPU模型学习代码的协作笔记本电脑 。
- 带有Keras TPU培训示例的协作笔记本，用于识别Fashion MNIST的衣服和鞋子 。
- Google Cloud中的Tensor处理器 。
- 张量处理器的体系结构和使用特点 。