Keras同时用多张显卡训练网络

References

官方文档:multi_gpu_model(https://keras.io/utils/#multi_gpu_model)以及Google。

误区

目前Keras是支持了多个GPU同时训练网络,非常容易,但是靠以下这个代码是不行的。

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

当你监视GPU的使用情况(nvidia-smi -l 1)的时候会发现,尽管GPU不空闲,实质上只有一个GPU在跑,其他的就是闲置的占用状态,也就是说,如果你的电脑里面有多张显卡,无论有没有上面的代码,Keras都会默认的去占用所有能检测到的GPU。

这行代码在你只需要一个GPU的时候时候用的,也就是可以让Keras检测不到电脑里其他的GPU。假设你一共有三张显卡,每个显卡都是有自己的标号的(0, 1, 2),为了不影响别人的使用,你只用其中一个,比如用gpu=1的这张,那么

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

然后再监视GPU的使用情况(nvidia-smi -l 1),确实只有一个被占用,其他都是空闲状态。所以这是一个Keras使用多显卡的误区,它并不能同时利用多个GPU。

目的

为什么要同时用多个GPU来训练?

单个显卡内存太小 -> batch size无法设的比较大,有时甚至batch_size=1都内存溢出(OUT OF MEMORY)

从我跑深度网络的经验来看,batch_size设的大一点会比较好,相当于每次反向传播更新权重,网络都可以看到更多的样本,从而不会每次iteration都过拟合到不同的地方去Don't Decay the Learning Rate, Increase the Batch Size。当然,我也看过有论文说也不能设的过大,原因不明... 反正我也没有机会试过。我建议的batch_size大概就是64~256的范围内,都没什么大问题。

但是随着现在网络的深度越来越深,对于GPU的内存要求也越来越大,很多入门的新人最大的问题往往不是代码,而是从Github里面抄下来的代码自己的GPU太渣,实现不了,只能降低batch_size,最后训练不出那种效果。

解决方案两个:一是买一个超级牛逼的GPU,内存巨大无比;二是买多个一般般的GPU,一起用。

第一个方案不行,因为目前即便最好的NVIDIA显卡,内存也不过十几个G了不起了,网络一深也挂,并且买一个牛逼显卡的性价比不高。所以、学会在Keras下用多个GPU是比较靠谱的选择。

实现

非常简洁

from model import unet G = 3 # 同时使用3个GPUwith tf.device("/gpu:0"):

M = unet(input_rows, input_cols, 1) model = keras.utils.training_utils.multi_gpu_model(M, gpus=G) model.compile(optimizer=Adam(lr=1e-5), loss='binary_crossentropy', metrics = ['accuracy']) model.fit(X_train, y_train, batch_size=batch_size*G, epochs=nb_epoch, verbose=0, shuffle=True, validation_data=(X_valid, y_valid)) model.save_weights('/path/to/save/model.h5')

问题

3.1 Compile the model

如果是普通的网络结构,那么没有问题,像上述的编译代码即可(model.compile(optimizer=Adam(lr=1e-5), loss='binary_crossentropy', metrics = ['accuracy'])) 。不过,如果是Multi-task的网络,例如Faster-RCNN,它由多个输出支路,也就是多个loss,在网络定义的时候一般会给命名,然后编译的时候找到不同支路layer的名字即可,就像这样:

model.compile(optimizer=optimizer, loss={'main_output': jaccard_distance_loss, 'aux_output': 'binary_crossentropy'}, metrics={'main_output': jaccard_distance_loss, 'aux_output': 'acc'},

loss_weights={'main_output': 1., 'aux_output': 0.5})

其中main_output和aux_output就是认为定义的layer name,但是如果用了keras.utils.training_utils.multi_gpu_model()以后,名字就自动换掉了,变成默认的concatenate_1, concatenate_2等等,因此你需要先model.summary()一下,打印出来网络结构,然后弄明白哪个输出代表哪个支路,然后重新编译网络,如下:

from keras.optimizers import Adam, RMSprop, SGD model.compile(optimizer=RMSprop(lr=0.045, rho=0.9, epsilon=1.0), loss={'concatenate_1': jaccard_distance_loss, 'concatenate_2': 'binary_crossentropy'}, metrics={'concatenate_1': jaccard_distance_loss, 'concatenate_2': 'acc'},

loss_weights={'concatenate_1': 1., 'concatenate_2': 0.5})

3.2 save the model

用多个GPU训练的模型有一个问题Keras没有解决,就是model.save()保存的时候报错

TypeError: can't pickle module objects

或是

RuntimeError: Unable to create attribute (object header message is too large)

原因是:

In https://keras.io/utils/#multi_gpu_model it clearly stated that the model can be used like the normal model, but it cannot be saved, very funny. I can't even perform reinforced training just because I cannot save the previous model trained with multiple GPUs. If trained with single GPU, the rest of my invested GPUs will become useless. Please urge the developer to look into this bug ASAP.

正常情况下Keras给你提供了自动保存最好的网络的函数(keras.callbacks.ModelCheckpoint()),它的内部是用model.save()来保存的,所以不能用了,你需要自己设计函数CustomModelCheckpoint()来保存最好的模型。

class CustomModelCheckpoint(keras.callbacks.Callback): def __init__(self, model, path):

self.model = model

self.path = path

self.best_loss = np.inf

def on_epoch_end(self, epoch, logs=None): val_loss = logs['val_loss']

if val_loss < self.best_loss: print("\nValidation loss decreased from {} to {}, saving model".format(self.best_loss, val_loss))

self.model.save_weights(self.path, overwrite=True)

self.best_loss = val_loss model.fit(X_train, y_train, batch_size=batch_size*G, epochs=nb_epoch, verbose=0, shuffle=True, validation_data=(X_valid, y_valid), callbacks=[CustomModelCheckpoint(model, '/path/to/save/model.h5')])

即便如此,要是模型还是太大,就需要下面的方法了,保存成npy格式而不是hdf5格式。

RuntimeError: Unable to create attribute (Object header message is too large)

  • model.get_weights(): returns a list of all weight tensors in the model, as Numpy arrays.
  • model.set_weights(weights): sets the values of the weights of the model, from a list of Numpy arrays. The arrays in the list should have the same shape as those returned by get_weights().

# save model

weight = self.model.get_weights() np.save(self.path+'.npy', weight)

# load model

weight = np.load(load_path) model.set_weights(weight)

3.3 Load the model

同样道理,当读入用多个显卡一起训练的网络文件.h的时候,也会报错

ValueError: You are trying to load a weight file containing 3 layers into a model with 1 layers.

原因是.h内部和单个GPU训练的存储不太一样,因此在读的时候也需要套一下keras.utils.training_utils.multi_gpu_model()这个函数。

from model import unetwith tf.device("/cpu:0"): M = unet(input_rows, input_cols, 1) model = keras.utils.training_utils.multi_gpu_model(M, gpus=G) model.load_weights(load_path)

然后就没问题啦。

原文发布于微信公众号 - 人工智能LeadAI(atleadai)

原文发表时间:2018-03-06

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏新智元

DeepMind 开源内部深度学习框架 Sonnet,研究通过梯度下降学习

【新智元导读】DeepMind 今天开源了最新的深度学习框架 Sonnet。Sonnet 被专门设计用于与 TensorFlow 协同工作,能够更方便、直接地构...

2793
来自专栏https://www.cnblogs.com/L

【自然语言处理篇】--Chatterbot聊天机器人

ChatterBot是一个基于机器学习的聊天机器人引擎,构建在python上,主要特点是可以自可以从已有的对话中进行学(jiyi)习(pipei)。

1475
来自专栏何俊林

深入浅出理解视频编码H.264结构

作者|Abson ? 作者介绍 Abson, 目前就职4399,系iOS开发工程师,主攻音视频开发技术。Abson的个人博客:http://simplecode...

2887
来自专栏量子位

看文本知语义:谷歌推一步到位自然语言理解框架SLING | 论文+代码

夏乙 编译自 Google Research Blog 量子位 出品 | 公众号 QbitAI 大多数自然语言理解(NLU)系统分析语言的过程是一条由分析步骤组...

3379
来自专栏AI科技评论

开发 | 如何为TensorFlow和PyTorch自动选择空闲GPU,解决抢卡争端

AI科技评论按:本文作者天清,原文载于其知乎专栏 世界那么大我想写代码,AI科技评论获授权发布。 项目地址:https://github.com/Quantum...

3648
来自专栏每日一篇技术文章

学习音视频解码你应该知道的东西

查看详细的视频编码介绍请访问视频编码 我们重点研究一下 H.26X 系列 特点:侧重网络传输 包括:H.261、H.262、H.263、H.263+、H....

942
来自专栏苍云横渡学习笔记

【day 11】python编程:从入门到实践学习笔记-基于Django框架的Web开发-Django入门(二)

学习笔记目录 【day 1】python编程:从入门到实践学习笔记-安装、变量和简单数据类型 【day 2】python编程:从入门到实践学习笔记-列表以及其操...

3017
来自专栏专知

【最新TensorFlow1.4.0教程01】TF1.4.0介绍与动态图机制 Eager Execution使用

【导读】主题链路知识是我们专知的核心功能之一,为用户提供AI领域系统性的知识学习服务,一站式学习人工智能的知识,包含人工智能( 机器学习、自然语言处理、计算机视...

3498
来自专栏机器之心

资源 | TensorFlow极简教程:创建、保存和恢复机器学习模型

选自Github 机器之心编译 参与:Jane W、李泽南 TensorFlow 是一个由谷歌发布的机器学习框架,在这篇文章中,我们将阐述 TensorFlow...

2537
来自专栏刘望舒

Android 人脸识别之人脸注册

2932

扫码关注云+社区