Colab提供了免费TPU,机器之心帮你试了试

最近机器之心发现谷歌的 Colab 已经支持使用免费的 TPU,这是继免费 GPU 之后又一重要的计算资源。我们发现目前很少有博客或 Reddit 论坛讨论这一点,而且谷歌也没有通过博客或其它方式做宣传。因此我们尝试使用该 TPU 训练简单的卷积神经网络,并对比它的运行速度。

我们在网上只发现比较少的信息与资源,最开始介绍 Colab 免费 TPU 的内容还是谷歌开发者 Sam Wittevee 最近的演讲 PPT。因此本文的测试和探索都是基于官方文档和实例所实现的,还有很多 TPU 特性没有考虑到,感兴趣的读者可查阅文末的参考资料,了解更多 Colab 免费 TPU 的特性。

本文所有的测试代码与结果都可以访问:https://colab.research.google.com/drive/1DpUCBm58fruGNRtQL_DiSVbT90spdZgm

试验 Colab 免费 TPU

首先我们需要确保 Colab 笔记本中运行时类型选择的是 TPU,同时分配了 TPU 资源。因此依次选择菜单栏中的「runtime」和「change runtime type」就能弹出以下对话框:

为了确保 Colab 给我们分配了 TPU 计算资源,我们可以运行以下测试代码。如果输出 ERROR 项,则表示目前的运行时并没有调整到 TPU,如果输出 TPU 地址及 TPU 设备列表,则表示 Colab 已经为我们分配了 TPU 计算资源。

如果查看以下测试代码的正常输出,Colab 会为「TPU 运行时」分配 CPU 和 TPU,其中分配的 TPU 工作站有八个核心,因此在后面配置的 TPU 策略会选择 8 条并行 shards。

import os
import pprint
import tensorflow as tf

if 'COLAB_TPU_ADDR' not in os.environ:
  print('ERROR: Not connected to a TPU runtime')
else:
  tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  print ('TPU address is', tpu_address)

  with tf.Session(tpu_address) as session:
    devices = session.list_devices()

  print('TPU devices:')
  pprint.pprint(devices)

目前,Colab 一共支持三种运行时,即 CPU、GPU(K80)和 TPU(据说是 TPU v2)。但我们不太了解 Colab 中的 GPU 和 TPU 在深度模型中的表现如何,当然后面会用具体的任务去测试,不过现在我们可以先用相同的运算试试它们的效果。因此我们首先尝试用简单的卷积运算测试它们的迭代时间。

在测试不同的硬件时,需要切换到不同的运行时。如下先定义 128 张随机生成的 256×256 图像,然后定义 256 个 5×5 的卷积核后就能执行卷积运算,其中魔术函数 %timeit 会自动多次执行,以产生一个更为精确的平均执行时间。

import tensorflow as tf 
import numpy as np
import timeit

tf.reset_default_graph()
img = np.random.randn(128, 256, 256, 3).astype(np.float32)
w = np.random.randn(5, 5, 3, 256).astype(np.float32)
conv = tf.nn.conv2d(img, w, [1,2,2,1], padding='SAME')

with tf.Session() as sess:
  # with tf.device("/gpu:0") as dev:
  %timeit sess.run(conv)

然而,是我们想当然了,使用 TPU 执行运算似乎需要特定的函数与运算,它不像 CPU 和 GPU 那样可以共用相同的代码。分别选择 CPU、GPU 和 TPU 作为运行时状态,运行上面的代码并迭代一次所需要的时间分别为:2.44 s、280 ms、2.47 s。从这里看来,仅修改运行时状态,并不会真正调用 TPU 资源,真正实现运算的还是 CPU。随后我们发现 TF 存在一个神奇的类 tf.contrib.tpu,似乎真正调用 TPU 资源必须使用它改写模型。

因此,根据文档与调用示例,我们将上面的卷积测试代码改为了以下形式,并成功地调用了 TPU。此外,因为每次都需要重新连接不同的运行时,所以这里的代码都保留了库的导入。虽然代码不太一样,但直觉上它的计算量应该和上面的代码相同,因此大致上能判断 Colab 提供的 GPU、TPU 速度对比。

import tensorflow as tf 
import numpy as np
import timeit
import os

tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']

tf.reset_default_graph()
def conv_op():
  img =  np.random.randn(128, 256, 256, 3).astype(np.float32)
  conv_w = np.random.randn(5, 5, 3, 256).astype(np.float32)
  conv = tf.nn.conv2d(img, conv_w, [1,2,2,1], padding='SAME')

tpu_ops = tf.contrib.tpu.batch_parallel(conv_op, [], num_shards=8)

with tf.Session(tpu_address) as sess:
  sess.run(tf.contrib.tpu.initialize_system())
  sess.run(tpu_ops)
  %timeit sess.run(tpu_ops)
  sess.run(tf.contrib.tpu.shutdown_system())

运行后出现了非常意外的结果,这样的卷积运算每一次迭代只需要 1.22 ms。如下图所示,很可能存在变量缓存等其它因素造成了一定程度的缓慢,但 TPU 的速度无可置疑地快。因此如果在 Colab 上测试模型,我们就更希望使用免费的 TPU,不过使用 TPU 需要改模型代码,这又比较麻烦。

尽管简单的卷积运算 TPU 要比 K80 快很多,但这只能给我们一个大致的猜想,因此我们需要测试完整的模型。注意在 tf.contrib.tpu 类中,它还提供了两种使用 TPU 的简单方法,即直接使用 Keras 接口和使用 TPUEstimator 构建模型。

在 tf.contrib.tpu 的文档中,我们发现 tf.contrib.tpu.keras_to_tpu_model 方法可以直接将 Keras 模型与对应的权重复制到 TPU,并返回 TPU 模型。该方法在输入 Keras 模型和在多个 TPU 核心上的训练策略后,能输出一个 Keras TPU 模型的实例,且可分配到 TPU 进行运算。

除此之外,另外一种调用 TPU 计算资源的方法是 tf.contrib.tpu.TPUEstimator,对于修正我们原来的TensorFlow 模型以适用 TPU,它可能是一种更方便的方式。根据文档所示,TPUEstimator 类继承自 Estimator 类,因此它不仅支持在 TPU 上运算,同时还支持 CPU 和 GPU 的运算。TPUEstimator 隐藏了非常多在 TPU 上训练的细节,例如为多个 TPU 核心复制多个输入和模型等。

TPU 调用文档地址:https://www.tensorflow.org/api_docs/python/tf/contrib/tpu

对比 TPU 与 GPU 的计算速度

为了简单起见,这里仅使用 Fashion-MNIST 数据集与简单的 5 层 卷积神经网络 测试不同的芯片性能。这个模型是基于 Keras 构建的,因为除了模型转换与编译,Keras 模型在 TPU 和 GPU 的训练代码都是一样的,且用 Keras 模型做展示也非常简洁。

几天前谷歌 Colab 团队发了一版使用 Keras 调用 TPU 的教程,因此我们就借助它测试 TPU 的训练速度。对于 GPU 的测试,我们可以修改该模型的编译与拟合部分,并调用 GPU 进行训练。所以整个训练的数据获取、模型结构、超参数都是一样的,不一样的只是硬件。

教程地址:https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/fashion_mnist.ipynb

以下是整个测试的公共部分,包含了训练数据的获取和模型架构。Keras 的模型代码非常好理解,如下第一个卷积层首先采用了批归一化,然后用 64 个 5×5 的卷积核实现卷积运算,注意这里采用的激活函数都是指数线性单元(ELU)。随后对卷积结果做 2×2 的最大池化,并加上一个随机丢弃率为 0.25 的Dropout层,最后得出的结果就是第一个卷积层的输出。

import tensorflow as tf
import numpy as np
import timeit

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

model = tf.keras.models.Sequential()

# 以下为第一个卷积层
model.add(tf.keras.layers.<mark data-type="technologies" data-id="e3995dcb-2c96-4671-b7ed-d5f2abae3c08">Batch</mark><mark data-type="technologies" data-id="a815978e-ab05-479e-b6b4-30176b3d5427">Normalization</mark>(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.Max<mark data-type="technologies" data-id="0a4cedf0-0ee0-4406-946e-2877950da91d">Pooling</mark>2D(pool_size=(2, 2), strides=(2,2)))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.25))

model.add(tf.keras.layers.<mark data-type="technologies" data-id="e3995dcb-2c96-4671-b7ed-d5f2abae3c08">Batch</mark><mark data-type="technologies" data-id="a815978e-ab05-479e-b6b4-30176b3d5427">Normalization</mark>(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.Max<mark data-type="technologies" data-id="0a4cedf0-0ee0-4406-946e-2877950da91d">Pooling</mark>2D(pool_size=(2, 2)))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.25))

model.add(tf.keras.layers.<mark data-type="technologies" data-id="e3995dcb-2c96-4671-b7ed-d5f2abae3c08">Batch</mark><mark data-type="technologies" data-id="a815978e-ab05-479e-b6b4-30176b3d5427">Normalization</mark>(input_shape=x_train.shape[1:]))
model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))
model.add(tf.keras.layers.Max<mark data-type="technologies" data-id="0a4cedf0-0ee0-4406-946e-2877950da91d">Pooling</mark>2D(pool_size=(2, 2), strides=(2,2)))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.25))

model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(256))
model.add(tf.keras.layers.Activation('elu'))
model.add(tf.keras.layers.<mark data-type="technologies" data-id="1c91194a-1732-4fb3-90c9-e0135c69027e">Dropout</mark>(0.5))
model.add(tf.keras.layers.Dense(10))
model.add(tf.keras.layers.Activation('softmax'))
model.summary()

在定义模型后,TPU 需要转化模型与编译模型。如下所示,keras_to_tpu_model 方法需要输入正常 Keras 模型及其在 TPU 上的分布式策略,这可以视为「TPU 版」的模型。完成模型的转换后,只需要像一般 Keras 模型那样执行编译并拟合数据就可以了。

注意两个模型的超参数,如学习率、批量大小和 Epoch 数量等都设置为相同的数值,且损失函数和最优化器等也采用相同的方法。

import os
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)
tpu_model.compile(
    optimizer=tf.train.Adam<mark data-type="technologies" data-id="fa50298e-1a85-4af0-ae96-a82708f4b610">Optimizer</mark>(learning_rate=1e-3, ),
    loss=tf.keras.losses.sparse_categorical_crossentropy,
    metrics=['sparse_categorical_accuracy']
)


def train_gen(batch_size):
  while True:
    offset = np.random.randint(0, x_train.shape[0] - batch_size)
    yield x_train[offset:offset+batch_size], y_train[offset:offset + batch_size]


%time tpu_model.fit_generator(train_gen(1024), epochs=5, steps_per_epoch=100, validation_data=(x_test, y_test))

最后在使用 GPU 训练模型时,我们会删除模型转换步骤,并保留相同的编译和拟合部分。训练的结果如下所示,Colab 提供的 TPU 要比 GPU 快 3 倍左右,一般 TPU 训练 5 个 Epoch 只需要 40 多秒,而 GPU 需要 2 分多钟。

Colab 使用免费 TPU 训练的信息摘要。

Colab 使用免费 GPU 训练的信息摘要。

最后,Colab 确实提供了非常强劲的免费 TPU,而且使用 Keras 或 TPUEstimator 也很容易重新搭建或转换已有的TensorFlow 模型。机器之心只是简单地试用了 Colab 免费 TPU,还有很多特性有待读者的测试,例如支持 TPU 的 PyTorch 1.0 或循环神经网络在 TPU 上的性能等。

参考资料:

  • 文档:https://www.tensorflow.org/api_docs/python/tf/contrib/tpu

  • 官方示例(Keras):https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb

  • 官方示例(TPUEstimator):https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpuestimator.ipynb

  • Sam Wittevee PPT:https://www.dropbox.com/s/jg7j07unw94wbom/TensorFlow%20Keras%20Colab%20TPUs.pdf?dl=0

  • Ceshine Lee 博客:https://medium.com/the-artificial-impostor/keras-for-tpus-on-google-colaboratory-free-7c00961fed69

工程 TPU 谷歌 Colab 计算资源

相关数据

激活函数 技术

Activation function

在 计算网络中, 一个节点的激活函数定义了该节点在给定的输入或输入的集合下的输出。标准的计算机芯片电路可以看作是根据输入得到"开"(1)或"关"(0)输出的数字网络激活函数。这与神经网络中的线性感知机的行为类似。 一种函数(例如 ReLU 或 S 型函数),用于对上一层的所有输入求加权和,然后生成一个输出值(通常为非线性值),并将其传递给下一层。

来源: 维基百科 Google ML glossary

神经网络 技术

Neural Network

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

来源:机器之心

Dropout 技术

Dropout

神经网络训练中防止过拟合的一种技术

来源:ImageNet Classification with Deep Convolutional

Exponential linear unit

一种激活函数, ELUs是对ReLU激活函数的一种演变,将激活函数更能够保持一个noise-robust状态。所以提出一个具有负值的激活函数,这可以使得平均激活接近于零,但它会以更小的参数饱和为负值的激活函数ELUs。

来源: Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)

Convolutional neural network

卷积神经网路(Convolutional Neural Network, CNN)是一种前馈神经网络,它的人工神经元可以响应一部分覆盖范围内的周围单元,对于大型图像处理有出色表现。卷积神经网路由一个或多个卷积层和顶端的全连通层(对应经典的神经网路)组成,同时也包括关联权重和池化层(pooling layer)。这一结构使得卷积神经网路能够利用输入数据的二维结构。与其他深度学习结构相比,卷积神经网路在图像和语音识别方面能够给出更好的结果。这一模型也可以使用反向传播算法进行训练。相比较其他深度、前馈神经网路,卷积神经网路需要考量的参数更少,使之成为一种颇具吸引力的深度学习结构。 卷积网络是一种专门用于处理具有已知的、网格状拓扑的数据的神经网络。例如时间序列数据,它可以被认为是以一定时间间隔采样的一维网格,又如图像数据,其可以被认为是二维像素网格。

来源:Goodfellow, I.; Bengio Y.; Courville A. (2016). Deep Learning. MIT Press. 维基百科

超参数 技术

Hyperparameter

在机器学习中,超参数是在学习过程开始之前设置其值的参数。 相反,其他参数的值是通过训练得出的。 不同的模型训练算法需要不同的超参数,一些简单的算法(如普通最小二乘回归)不需要。 给定这些超参数,训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式,并且必须对其进行调整以便模型能够最优地解决机器学习问题。 在实际应用中一般需要对超参数进行优化,以找到一个超参数元组(tuple),由这些超参数元组形成一个最优化模型,该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源: Wikipedia

学习率 技术

Learning rate

在使用不同优化器(例如随机梯度下降,Adam)神经网络相关训练中,学习速率作为一个超参数控制了权重更新的幅度,以及训练的速度和精度。学习速率太大容易导致目标(代价)函数波动较大从而难以找到最优,而弱学习速率设置太小,则会导致收敛过慢耗时太长

来源:Liu, T. Y. (2009). Learning to rank for information retrieval. Foundations and Trends® in Information Retrieval, 3(3), 225-331. Wikipedia

损失函数 技术

Loss function

在数学优化,统计学,计量经济学,决策理论,机器学习和计算神经科学等领域,损失函数或成本函数是将一或多个变量的一个事件或值映射为可以直观地表示某种与之相关“成本”的实数的函数。

来源: Wikipedia

规范化 技术

Normalization

规范化:将属性数据按比例缩放,使之落入一个小的特定区间,如-1.0 到1.0 或0.0 到1.0。 通过将属性数据按比例缩放,使之落入一个小的特定区间,如0.0到1.0,对属性规范化。对于距离度量分类算法,如涉及神经网络或诸如最临近分类和聚类的分类算法,规范化特别有用。如果使用神经网络后向传播算法进行分类挖掘,对于训练样本属性输入值规范化将有助于加快学习阶段的速度。对于基于距离的方法,规范化可以帮助防止具有较大初始值域的属性与具有较小初始值域的属相相比,权重过大。有许多数据规范化的方法,包括最小-最大规范化、z-score规范化和按小数定标规范化。

来源:Jiawei Han;Micheline Kamber著 数据挖掘概念与技术 机械工业出版社

优化器 技术

Optimizer

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

来源: 维基百科

池化 技术

Pooling

池化(Pooling)是卷积神经网络中的一个重要的概念,它实际上是一种形式的降采样。有多种不同形式的非线性池化函数,而其中“最大池化(Max pooling)”是最为常见的。它是将输入的图像划分为若干个矩形区域,对每个子区域输出最大值。直觉上,这种机制能够有效的原因在于,在发现一个特征之后,它的精确位置远不及它和其他特征的相对位置的关系重要。池化层会不断地减小数据的空间大小,因此参数的数量和计算量也会下降,这在一定程度上也控制了过拟合。通常来说,CNN的卷积层之间都会周期性地插入池化层。

来源: cs231n

权重 技术

Weight

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

来源:Google AI Glossary

批次 技术

Batch

模型训练的一次迭代(即一次梯度更新)中使用的样本集。

来源: Google ML glossary

张量 技术

Tensor

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

来源: 维基百科

TensorFlow 技术

TensorFlow

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

来源: 维基百科

思源

机器之心编辑

我来评几句
登录后评论

已发表评论数()

相关站点

+订阅
热门文章