利用Colab上的TPU训练Keras模型（完整版）

最新推荐文章于 2024-03-09 11:52:10 发布

big91987

最新推荐文章于 2024-03-09 11:52:10 发布

阅读量1.6w

点赞数 17

分类专栏：深度学习工具文章标签：深度学习工具

本文链接：https://blog.csdn.net/big91987/article/details/87898100

版权

深度学习工具专栏收录该内容

1 篇文章 0 订阅

订阅专栏

0. 前言

Colab是由Google提供的云计算服务，通过它可以让开发者很方便的使用google的免费资源（CPU、GPU、TPU）来训练自己的模型；网络上已经有很多相关的资料了，但是都比较零散，因此我整理了一些资料，并记录了如何从0开始，一步一步的使用colab在TPU模式下训练keras的模型。（需要读者有一定的keras和notebook的基础。）

以下是本文没有覆盖的内容：
- 如何使用keras建立和训练模型
- 如何搭梯子

可以直接访问我的共享colab，查看相应步骤：
https://colab.research.google.com/drive/1hq9ODv77B2cUkpeI7XM_W7bIsko6vWQH

1. Colab的建立

1.1 登录google云盘

登陆谷歌云盘：driver.google.com
需要梯子及google账号，登录进去以后，在右边的空白处，可以通过右键新建文件或者文件夹
在这里插入图片描述

1.2 建立colaboratory文件

通过右键可以建立文件夹、文件、上传下载，
首先通过右键/更多/Colaboratory新建一个jpynb文件；
然后建立‘for_colab_cache’的文件夹（名字随便取），后面用来挂载到colab上，用来保存训练出来的模型
在这里插入图片描述
如果在更多里面没有找到“Colaboratory”，可以点击“+关联更多应用”，然后在里面搜索并添加colaboratory

1.3 运行colaboratory

右键该文件，选择打开方式“colaboratory”，也可以双击文件，然后选择打开方式为“colaboratory”
在这里插入图片描述

2. Colab的配置

打开文件以后，会跳转到colab网站上对应刚才新建的notebook上，网址是类似
https://colab.research.google.com/drive/XXXXXXXX 的格式
显示的是一个类似notebook的界面：
在这里插入图片描述

2.1 查看RAM与磁盘

注意右上方的RAM和硬盘，鼠标放到上面可以看到相应的信息：
在这里插入图片描述

2.2 切换服务器模式（CPU/GPU/TPU）

点击修改/笔记本设置，可以切换服务器模式，包括使用的python版本，和硬件型号：
在这里插入图片描述
当模式被切换以后，服务的后端会自动切换到一个对应的类型的新服务器上，因此之前安装的软件以及本地生成的文件会被清空，需要重新设置，右上方会变成“正在初始化” “分配”之类的状态，当分配完毕后，该位置会显示新服务器上的RAM及磁盘使用情况
有时候可能会分配不到服务器，右边会变为“忙碌”状态。

2.3 设置笔记本显示行号

点击工具/偏好设置，在里面勾选“显示行号”，然后点击保存
在这里插入图片描述

2.4 在notebook中执行linux命令

在notebook里面，建立代码单元（点那个"+代码"）,在代码单元格中通过[!+shell命令]的格式，来运行linux命令，然后点击shift+enter来执行命令：
在这里插入图片描述
可以看到，notebook中的初始路径为‘/content’

3. 文件访问

由于colab的后端服务器是动态分配的，因此生成或者上传的文件在下次访问的时候就不在新服务器上了，因此涉及到文件访问问题
colab中提供了几种方式用于访问文件，比较常用的有上传/下载，以及挂载谷歌云盘

3.1 文件的上传/下载

在代码单元中执行如下命令上传文件：

from google.colab import files

uploaded = files.upload()

for fn in uploaded.keys():
    print('User uploaded file "{name}" with length {length} bytes'.format(name=fn, length=len(uploaded[fn])))

执行完毕后，点击“选择文件”按钮，上传相应文件：
在这里插入图片描述
在单元格中执行如下代码，下载相应的文件：

from google.colab import files

files.download('./weight.best.hdf5')

3.2 挂载driver中的文件夹

通过如下方式可以挂载谷歌云盘

# Load the Drive helper and mount
from google.colab import drive

# This will prompt for authorization.
drive.mount('/content/drive')

执行以后，需要点击单元格下方的链接在弹出的网页上进行认证，然后从中复制验证码，粘贴到文本框中，点击回车即可
在这里插入图片描述

云盘被挂载到/content/drive/'My Drive’中，可以通过软连接，将之前建立的’for_colab_cache’连接到 /content下，方便后续访问

4. 设置TPU加速

首先，按照上面的设置方法，将notebook设置为TPU模式，
然后修改自己的代码加入TPU初始化、模型转换的相关步骤

4.1 初始化TPU环境

# 准备TPU环境（可选）：
import os
# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)

4.2 将keras模型转化为TPU模型

以keras为例，假设已经定义好了一个模型model，需要将其转换为TPU类型的模型

model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

执行以后，可以看到TPU相应的log：
在这里插入图片描述

进行模型转换的时候可能会报错，提示model的类型不匹配，这有可能是因为你调用了标准keras的model，而TPU只兼容tensorflow.python.keras.XXX 格式的库，需要修改keras库的引用

4.3 使用tf优化器（可选）

目前标准TPU对keras标准的优化器（keras.optimizers）支持有限，有时候可能使用tf.train.的优化器进行编译，读者可根据需要选取，如果没有特殊需求，建议还是使用keras自己的优化器

model.compile(optimizer=tf.train.RMSPropOptimizer(learning_rate=1e-3), loss='categorical_crossentropy', metrics=['caccuracy'])

需要注意的是，keras的某些callback例如LearningRateScheduler，与tf的优化器不兼容，因为学习率的变量，在tf优化器中是以learning_rate保存的，而在keras的优化器中是用lr保存的，另外在训练的时候tf优化器会被拆成多个shard扔到不同的核心上，因此学习率的路径也不是固定的，我重写了一个callback类，暂时可以避免该问题，

# 重写learning callback适配tf.train优化器
class LearningRateScheduler(keras.callbacks.Callback):

    def __init__(self, schedule, verbose=0):
        super(LearningRateScheduler_zjz, self).__init__()
        self.schedule = schedule
        self.verbose = verbose
      
    def get_lr_opt(self):
        if not hasattr(self, 'model'):
            return None,None
        tmp = self.model
        lr_attr_name = None
        lr_name_list = ['lr', '_lr', '_learning_rate', 'learning_rate']
        opt_name_list = ['optimizer', '_opt']
        while True:
            for lr_name in lr_name_list:
                if hasattr(tmp, lr_name):
                    lr_attr_name = lr_name
                    break
            if lr_attr_name != None:
                return tmp, lr_attr_name
            for opt_name in opt_name_list:
                if hasattr(tmp, opt_name):
                    tmp = getattr(tmp, opt_name)
                    break
        return None,None

    def on_epoch_begin(self, epoch, logs=None):
        opt, lr_name =  self.get_lr_opt()
        if opt is None:
            raise ValueError('no attr _lr, _learning or lr found')
        try:  # new API
            lr = getattr(opt, lr_name)
            lr = self.schedule(epoch, lr)
        except TypeError:  # Support for old API for backward compatibility
            lr = self.schedule(epoch)
        if not isinstance(lr, (float, np.float32, np.float64)):
            raise ValueError('The output of the "schedule" function '
                         'should be float.')
        setattr(opt, lr_name, lr)
        if self.verbose > 0:
            print('\nEpoch %05d: LearningRateScheduler reducing learning '
              'rate to %s.' % (epoch + 1, lr))
      
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        opt, lr_name =  self.get_lr_opt()
        if opt is None:
            raise ValueError('no attr _lr, _learning or lr found')
        logs['lr'] = float(getattr(opt, lr_name))

4.4 设定fit时的batch_size

转换以后，需要注意云TPU默认给模型分了8个核心，每个核心跑一份batch，所以在fit的时候batch_size应为实际模型中batch_size的8倍，一般起点选择128*8，如果是较大的模型，可能显存会超，此时可以慢慢降低batch_size。

model.fit(x_train, y_train,epochs=100,batch_size=128 * 8)

fit以后可以看到相应的log：
在这里插入图片描述
以cifar10 resnetv1为例，在GPU上大约一个epoch 2分钟左右（batch_size=64），在TPU上平均25s一个epoch(batch_size=128)，可以看到提升还是很明显的：

如果读者使用keras的fit_generator方法进行训练，需要注意不要设置worker和use_multiprocessing，似乎多worker的时候和TPU不兼容，会导致训练一直卡在epoch 1时，感觉是keras死锁了；
我的一个典型调用方式时这样的：

model.fit_generator(datagen.flow(x_train, y_train, batch_size=128 * 8),
                    epochs=epochs,
                    steps_per_epoch=x_train.shape[0] // batch_size)

4.5 TPU模型的保存与恢复

虽然TPU训练速度很快，但是由于网络原因，可能出现网络中断重连的情况，此时后端服务器可能会被切换到一个新的TPU服务器上，因此涉及到模型的保存与恢复
保存之前需要先将其转换为CPU模型，可以保存在前面挂载的云盘的文件夹中：

save_model =  model.sync_to_cpu()
save_model.save('/content/for_colab_cache/best.hdf5')

读取模型以后，需要重新将其转换为TPU模型，可能需要重新编译优化器：

model = models.load_model('/content/for_colab_cache/best.hdf5')
model = tf.contrib.tpu.keras_to_tpu_model(
	model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
    	tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
model.compile(optimizer=tf.train.RMSPropOptimizer(learning_rate=1e-3), loss='categorical_crossentropy', metrics=['caccuracy'])

4.6 使用权重推理

可以直接调用keras的接口保存权重用于推理：

model.save_weights('./tpu.h5', overwrite=True)
new_model = make_model(batch_size=None)
new_model.load_weights('./tpu.h5')
new_model.summary()
new_model.evaluate(x_test, y_test)

4.7 Debug

（1）TPU上下文被破坏
出现如下提示，说明TPU环境上下文被破坏了，需要重新安装tensorflow：

InvalidArgumentError: No OpKernel was registered to support Op ‘ConfigureDistributedTPU’ used by {{node ConfigureDistributedTPU}}with these attrs: [embedding_config="", tpu_embedding_config="", is_global_init=false]
Registered devices: [CPU, XLA_CPU]

此时执行：!pip uninstall tensorflow 先删除tensorflow
再执行： !pip install tensorflow 重新安装tensorflow试试