4步在Google Colab的云TPU上训练模型

一路向北g
920 阅读7分钟

普通的TensorFlow模型，它的计算成本太高，无法在普通的工作笔记本电脑上训练。我明白，也经历过，老实说，看到我的笔记本在试图训练一个模型后连续两次崩溃是很痛苦的事情。

在这篇文章中，我将详细介绍如何使用谷歌Colab在云端的TPU上训练任何模型的步骤。相信我，在这之后，您再也不想碰笨重的 CPU 了。

TL;DR： 这篇文章向你展示了在TPU上训练任何TensorFlow模型是多么容易，只需对你的代码做很少的改动。

什么是TPU？

张量处理单元（TPU）是一个加速器--由谷歌大脑的硬件工程师定制--专门用于训练深度和计算量大的ML模型。

image.png

这就是TPU的样子。它有4个巨大的散热器，位于一个带有一些非常先进和强大的电路的电路板之上。

让我们客观地看待事物，让您了解 TPU 有多么出色和强大。标准的 MacBook Pro Intel CPU 每个时钟周期可以执行一些操作。一个标准的现成 GPU 每个周期可以执行数万次操作。最先进的 TPU 每个周期可以执行数十万次操作（有时高达 128K OPS）。

要了解规模，想象一下使用这些设备打印一本书。CPU 可以逐个字符地打印。GPU 一次可以打印几个字。热塑性聚氨酯？好吧，它一次可以打印一整页。这是我们现在可以使用的惊人速度和力量；感谢谷歌让每个人都能使用高性能硬件。

如果您对 TPU 的内部工作原理以及它的惊人之处感兴趣，请查看 Google Cloud 博客文章，他们在这里讨论了从硬件到软件的所有内容 。您可以找到有关机制、硬件规格、优势、限制、约束等的信息，以帮助您完成项目！

预处理MNIST

在本教程中，我将通过一个快速而简单的MNIST模型。请注意，这个模型可以是你想做的任何事情。为了更好地帮助你了解正在发生的事情，我选择了良好的旧MNIST（同样，你可以选择一个数据集）。

首先，让我们开始提取和预处理我们的数据集。这不应该是一个很大的问题：

从 tensorflow.keras.datasets导入 tensorflow 作为 tf从 tensorflow.keras.utils
导入 mnist
导入 to_categorical# 这可以是任何数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data() 
x_train = x_train.reshape([x_train.shape[0], 784]) 
x_test = x_test.reshape([x_test.shape [0], 784]) 
x_train = x_train / 255 
x_test = x_test / 255y_train = to_categorical(y_train, 10) 
y_test = to_categorical(y_test, 10)打印 (x_train.shape, y_train.shape)
打印 (x_test.shape, y_test.shape)# >> (60000, 784), (60000, 10) 
# >> (10000, 784), (10000, 10)

我已尝试使本教程尽可能全面，以便您可以以不可思议的速度训练您的模型，并感觉自己处于世界之巅。我们将采取 4 个步骤在 TPU 上训练我们的模型：

  1. 连接到可用的 TPU 实例
  2. 初始化分布式训练策略
  3. 根据上述策略建立我们的模型
  4. 训练模型，感觉自己像个巫师

让我们直接开始吧!

注意： 欢迎使用 tf.data ，对你的数据集进行预处理 ，或者将其转换为 TFRecords *，这些步骤不是必需的，但在您处理非常大的数据集的情况下可能有助于处理内存不足的情况。

连接到一个TPU

当我在Colab上捣鼓TPU的时候，连接到一个TPU是最乏味的。花了好几个小时在网上搜索和看教程，但我终于能够完成它了。

我们首先需要将本地运行的Colab笔记本连接到TPU的运行时间。要改变运行时间，只需点击导航栏中的Runtime 标签。将显示一个下拉菜单，你可以从中选择Change runtime type 选项。会出现一个弹出窗口，你可以从下拉选择器中选择TPU 选项。

image.png

弹出的窗口将看起来像这样。将运行时间从CPU切换到TPU。

为了连接到一个TPU实例，我们需要创建一个TPUClusterResolver ，它可以接收可用的设备并为我们提供一个：

Detect hardware 
try: 
tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection 
except ValueError: # If TPU not found                           
tpu = None

这应该会把我们当前的Colab会话连接到一个可用的TPU实例上。让我们继续初始化策略，最后调用TPU。

注意： 如果你在多次尝试和重新运行后仍难以连接到一个TPU实例，你可以将运行时类型改为 GPU 来代替。该代码仍然可以有效地编译和运行。

初始化分布式训练策略

现在我们已经改变了运行时并获得了一个可用的TPU列表，我们需要通过创建一个分布式训练策略来调用TPU--一个围绕模型的包装，使其与多核训练兼容。

Select appropriate distribution strategy 
if tpu:
tf.tpu.experimental.initialize_tpu_system(tpu)                 
strategy = tf.distribute.experimental.TPUStrategy(tpu, 
steps_per_run=128) print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])                         
else:
strategy = tf.distribute.get_strategy() # Default strategy that works on CPU and single GPU 
print('Running on CPU instead')                                 
print("Number of accelerators: ", strategy.num_replicas_in_sync)

这resolver使我们能够访问 TPU，以便我们最终可以在其上构建并行分布的管道。这是必要的步骤，因为 TPU 是分布式训练处理器，不像传统 CPU 那样是单核的。使用这种策略方法，跳上 TPU 非常简单！它应该给你一个可用实例的列表：

image.png

我们有8个设备供使用

在分布式训练策略下建立我们的模型

现在有了训练策略，我们可以继续使用该策略来建立我们的模型：

with strategy.scope(): 
"""
This essentailly takes our model and makes it                   
compatible to train on a TPU.                                   
"""                                                             
model = tf.keras.models.Sequential()                           
model.add(tf.keras.layers.Dense(512, input_size=[784,], activation='relu'))
model.add(tf.keras.layers.Dense(256, activation='relu'))       
model.add(tf.keras.layers.Dense(128, activation='relu'))        
model.add(tf.keras.layers.Dense(64, activation='relu'))         
model.add(tf.keras.layers.Dense(10, activation='softmax'))     
""" 
Compiling the model using the RMSProp optimizer                 
and Sparse Categorical Crossentropy loss                        
""" 
model.compile(                                                  
optimizer=tf.train.RMSPropOptimizer(learning_rate=1e-2),        
loss=tf.keras.losses.sparse_categorical_crossentropy,           
metrics=['sparse_categorical_accuracy']                          
) 
model.summary()

这看起来像很多专业术语，但它所做的只是把一个普通的tf.keras 模型（通常在CPU上运行）放在TPU上，并在所有可用的TPU核心上自动分配。

训练我们的模型

这是这个过程中令人兴奋的部分。我们终于可以在云TPU上免费训练我们的模型了。有了这么大的能量，让我们好好利用它，在MNIST上进行训练（我知道...很反常）。

history = model.fit(x_train,
y_train, 
epochs=20, 
steps_per_epoch=50 
model.save_weights('./mnist_model.h5', overwrite=True)

确保实例的数量完全可以被steps_per_epoch ，以便在
训练期间使用所有的实例。例如，我们的训练集里有60000个实例。60000
是可以被50整除的，所以这意味着我们所有的实例都被送入模型，没有任何遗漏。如果你点击运行，它应该在短时间内开始训练TPU实例：

image.png

这是我在Colab上建立的神经机器翻译模型的TPU训练周期。

最后的总结

就这样，训练应该很快开始了。Colab会启动一个TPU，并在上面上传模型架构。你应该很快就能在终端输出中看到经典的Keras进度条风格的布局。祝贺你!你已经成功地使用了一个TPU。

回想起来，我一直认为在TPU上训练模型是一件只有ML Wizards--那些有多年经验和技能的人--才能处理的事情。TPU训练一直是我无法掌握的一件大事，仅仅是因为围绕它的炒作（使它看起来非常难以使用）。在我最疯狂的梦想中，我从未想过它会像在我预先存在的模型中添加不到10行代码一样简单。

读完这篇文章后，我想让你知道，任何人都可以在加速的硬件上进行训练，无论经验如何。这只是谷歌人工智能为在大众中实现机器学习和人工智能的民主化所采取的众多措施之一。

avatar
文章
阅读
粉丝