[Keras] 使用Keras调用多GPU，并保存模型

官方文档对于如何调用多GPU已经说的很清楚：multi_gpu_model，但仍有些细节，值得探讨：

keras.utils.multi_gpu_model(model, gpus)

将模型在多个GPU上复制

特别地，该函数用于单机多卡的数据并行支持，它按照下面的方式工作：
（1）将模型的输入分为多个子batch
（2）在每个设备上调用各自的模型，对各自的数据集运行
（3）将结果连接为一个大的batch（在CPU上）
例如，你的batch_size是64而gpus=2，则输入会被分为两个大小为32的子batch，在两个GPU上分别运行，通过连接后返回大小为64的结果。 该函数线性的增加了训练速度，最高支持8卡并行。

*该函数只能在tf后端下使用

参数如下：

• model: Keras模型对象，为了避免OOM错误（内存不足），该模型应在CPU上构建，参考下面的例子。
• gpus: 大或等于2的整数，要并行的GPU数目。
  该函数返回Keras模型对象，它看起来跟普通的keras模型一样，但实际上分布在多个GPU上。

例子：

import tensorflow as tf
from keras.applications import Xception
from keras.utils import multi_gpu_model
import numpy as np

num_samples = 1000
height = 224
width = 224
num_classes = 1000

# Instantiate the base model
# (here, we do it on CPU, which is optional).
with tf.device('/cpu:0'):
    model = Xception(weights=None,
                     input_shape=(height, width, 3),
                     classes=num_classes)

# Replicates the model on 8 GPUs.
# This assumes that your machine has 8 available GPUs.
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(loss='categorical_crossentropy',
                       optimizer='rmsprop')

# Generate dummy data.
x = np.random.random((num_samples, height, width, 3))
y = np.random.random((num_samples, num_classes))

# This `fit` call will be distributed on 8 GPUs.
# Since the batch size is 256, each GPU will process 32 samples.
parallel_model.fit(x, y, epochs=20, batch_size=256)

但是在parallel_model.fit()结束后，使用代码parallel_model.save()保存却出现错误：

parallel_model.save('test.h5')
Traceback (most recent call last):

  File "", line 1, in 
    parallel_model.save('test.h5')

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/site-packages/keras/engine/topology.py", line 2556, in save
    save_model(self, filepath, overwrite, include_optimizer)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/site-packages/keras/models.py", line 107, in save_model
    'config': model.get_config()

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/site-packages/keras/engine/topology.py", line 2397, in get_config
    return copy.deepcopy(config)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 230, in _deepcopy_list
    y.append(deepcopy(a, memo))

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 237, in _deepcopy_tuple
    y.append(deepcopy(a, memo))

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 237, in _deepcopy_tuple
    y.append(deepcopy(a, memo))

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 190, in deepcopy
    y = _reconstruct(x, rv, 1, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 334, in _reconstruct
    state = deepcopy(state, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 190, in deepcopy
    y = _reconstruct(x, rv, 1, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 334, in _reconstruct
    state = deepcopy(state, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 264, in _deepcopy_method
    return type(x)(x.im_func, deepcopy(x.im_self, memo), x.im_class)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 190, in deepcopy
    y = _reconstruct(x, rv, 1, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 334, in _reconstruct
    state = deepcopy(state, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 230, in _deepcopy_list
    y.append(deepcopy(a, memo))

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 190, in deepcopy
    y = _reconstruct(x, rv, 1, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 334, in _reconstruct
    state = deepcopy(state, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 298, in _deepcopy_inst
    state = deepcopy(state, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 190, in deepcopy
    y = _reconstruct(x, rv, 1, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 334, in _reconstruct
    state = deepcopy(state, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 163, in deepcopy
    y = copier(x, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 257, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)

  File "/home/dexter/anaconda2/envs/tensorflow/lib/python2.7/copy.py", line 182, in deepcopy
    rv = reductor(2)

TypeError: can't pickle thread.lock objects

这个问题困扰了我很久，最后在 keras-team/keras/issues#8446&issues#8253找到正解。
不过当时提问者报错为：

TypeError: can’t pickle module objects

与我的TypeError: can't pickle thread.lock objects大同小异，解决方法如下：

意思就是直接使用传入方法 keras.utils.multi_gpu_model(model, gpus)中的 model即可，而不要使用返回的 parallel_model，即：

model.save('xxx.h5')