当多个GPU用于训练时,加载预训练模型失败

时间:2017-06-05 14:49:15

标签: neural-network deep-learning keras backend

我已经培训了一个网络模型并通过checkpoint = ModelCheckpoint(filepath='weights.hdf5')回调保存了它的权重和架构。在培训期间,我通过调用以下功能使用多个GPU:

def make_parallel(model, gpu_count):
    def get_slice(data, idx, parts):
        shape = tf.shape(data)
        size = tf.concat([ shape[:1] // parts, shape[1:] ],axis=0)
        stride = tf.concat([ shape[:1] // parts, shape[1:]*0 ],axis=0)
        start = stride * idx
        return tf.slice(data, start, size)

    outputs_all = []
    for i in range(len(model.outputs)):
        outputs_all.append([])

    #Place a copy of the model on each GPU, each getting a slice of the batch
    for i in range(gpu_count):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:

                inputs = []
                #Slice each input into a piece for processing on this GPU
                for x in model.inputs:
                    input_shape = tuple(x.get_shape().as_list())[1:]
                    slice_n = Lambda(get_slice, output_shape=input_shape, arguments={'idx':i,'parts':gpu_count})(x)
                    inputs.append(slice_n)                

                outputs = model(inputs)

                if not isinstance(outputs, list):
                    outputs = [outputs]

                #Save all the outputs for merging back together later
                for l in range(len(outputs)):
                    outputs_all[l].append(outputs[l])

    # merge outputs on CPU
    with tf.device('/cpu:0'):
        merged = []
        for outputs in outputs_all:
            merged.append(merge(outputs, mode='concat', concat_axis=0))

        return Model(input=model.inputs, output=merged)

使用测试代码:

from keras.models import Model, load_model
import numpy as np
import tensorflow as tf

model = load_model('cpm_log/deneme.hdf5')

x_test = np.random.randint(0, 255, (1, 368, 368, 3))

output = model.predict(x = x_test, batch_size=1)

print output[4].shape

我收到了以下错误:

Traceback (most recent call last):
  File "cpm_test.py", line 5, in <module>
    model = load_model('cpm_log/Jun5_1000/deneme.hdf5')
  File "/usr/local/lib/python2.7/dist-packages/keras/models.py", line 240, in load_model
    model = model_from_config(model_config, custom_objects=custom_objects)
  File "/usr/local/lib/python2.7/dist-packages/keras/models.py", line 301, in model_from_config
    return layer_module.deserialize(config, custom_objects=custom_objects)
  File "/usr/local/lib/python2.7/dist-packages/keras/layers/__init__.py", line 46, in deserialize
    printable_module_name='layer')
  File "/usr/local/lib/python2.7/dist-packages/keras/utils/generic_utils.py", line 140, in deserialize_keras_object
    list(custom_objects.items())))
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/topology.py", line 2378, in from_config
    process_layer(layer_data)
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/topology.py", line 2373, in process_layer
    layer(input_tensors[0], **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/topology.py", line 578, in __call__
    output = self.call(inputs, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/keras/layers/core.py", line 659, in call
    return self.function(inputs, **arguments)
  File "/home/muhammed/DEV_LIBS/developments/mocap/pose_estimation/training/cpm/multi_gpu.py", line 12, in get_slice
    def get_slice(data, idx, parts):
NameError: global name 'tf' is not defined

通过检查错误输出,我决定问题出在并行化代码上。但是,我无法解决问题。

1 个答案:

答案 0 :(得分:1)

您可能需要使用custom_objects来启用模型加载。

import tensorflow as tf
model = load_model('model.h5', custom_objects={'tf': tf,})