在 colab 中使用 keras_to_tpu_model 时,TPU 运行速度与 CPU 一样慢

Posted

技术标签:

【中文标题】在 colab 中使用 keras_to_tpu_model 时,TPU 运行速度与 CPU 一样慢【英文标题】:TPU runs as slow as CPU when using keras_to_tpu_model in colab 【发布时间】:2019-04-18 02:43:34 【问题描述】:

我使用tf.contrib.tpu.keras_to_tpu_model 使我的代码能够在 TPU 上运行,但是完成一个 epoch 需要 170 小时,而 CPU 需要相同的时间,GPU 每个 epoch 只需要 40 小时。我尝试调整批量大小,但是没有任何改变。而且我测试过输入函数在 GPU 上运行时可能会占用 20% 的运行时间,所以我认为这可能不是主要原因。

这是我的代码:https://github.com/WangHexie/DHNE/blob/master/src/hypergraph_embedding.py

在 colab 上运行:

    TPU:https://colab.research.google.com/gist/WangHexie/30c385509f9cd93be747f04c39f039a4/tpu-error.ipynb GPU:https://colab.research.google.com/gist/WangHexie/5bfac53bf92ef0ad527f15ddbf8705e1/-gpu-ipynb.ipynb

型号:

def build_model(self):
    self.inputs = [Input(shape=(self.options.dim_feature[i], ), name='input_'.format(i), dtype='float') for i in range(3)]

    self.encodeds = [Dense(self.options.embedding_size[i], activation='tanh', name='encode_'.format(i))(self.inputs[i]) for i in range(3)]
    self.decodeds = [Dense(self.options.dim_feature[i], activation='sigmoid', name='decode_'.format(i),
                    activity_regularizer = regularizers.l2(0.0))(self.encodeds[i]) for i in range(3)]

    self.merged = concatenate(self.encodeds, axis=1)
    self.hidden_layer = Dense(self.options.hidden_size, activation='tanh', name='full_connected_layer')(self.merged)
    self.ouput_layer = Dense(1, activation='sigmoid', name='classify_layer')(self.hidden_layer)

    self.model = Model(inputs=self.inputs, outputs=self.decodeds+[self.ouput_layer])

    self.model.compile(optimizer=tf.train.AdamOptimizer(learning_rate=self.options.learning_rate),
            loss=[self.sparse_autoencoder_error]*3+['binary_crossentropy'],
                          loss_weights=[self.options.alpha]*3+[1.0],
                          metrics=dict([('decode_'.format(i), 'mse') for i in range(3)]+[('classify_layer', 'accuracy')]))
    self.model = tf.contrib.tpu.keras_to_tpu_model(
        self.model,
        strategy=tf.contrib.tpu.TPUDistributionStrategy(
            tf.contrib.cluster_resolver.TPUClusterResolver(
                tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
        )
    )
    self.model.summary()

【问题讨论】:

【参考方案1】:

截至 2019 年 2 月 20 日,函数 tf.contrib.tpu.keras_to_tpu_model 已被弃用。因此,您应该重新尝试使用新的Distribution Strategy 函数转换您的模型。分布式训练的深度指南可以在here找到。

我还注意到您使用数据类型 float 作为输入值。在 CPython 中,默认位值为 64 位。目前,TPU 的功能最优化使用 16 位浮点数,因此您应该将输入减少到 8 位或 16 位。位值越低,模型的处理速度就越快。

因此还建议利用量化,将浮点权重转换为 8 位整数。量化训练有两种类型:post-training quantization 和quantization-aware training。

有关 Google Cloud Platform 上的 TPU 的更多信息,您可以参考Cloud TPU documentation,有关 TPU 系统架构的更多信息,您可以参考 Google 的 this 文档,因为它正确解释了 TPU 的设计方式。

【讨论】:

以上是关于在 colab 中使用 keras_to_tpu_model 时,TPU 运行速度与 CPU 一样慢的主要内容,如果未能解决你的问题,请参考以下文章

如何在 Google colab 中更改 Keras/tensorflow 版本?

Google colab 无法处理 hdf5 文件

为啥我的 Colab 会话内存不足?

colab在更换python包版本时,如何正确重启

colab在更换python包版本时,如何正确重启

在 Colab TPU 上运行 Pytorch 堆叠模型