如何将具有自定义 keras 层（.h5）的 keras 模型冻结到张量流图（.pb）？

Question

I am trying to implement a Faster-RCNN model for object detection written by Yinghan Xu . After I have trained and saved the model with model_all.save('filename.h5') , I am trying to freeze the Keras model as TensorFlow graph (as .pb ) for inference using keras_to_tensorflow.py written by Amir Abdi. But when I try to convert it, I get a ValueError: Unknown layer: roipoolingconv due to a custom RoiPoolingConv layer:

class RoiPoolingConv(Layer):
'''ROI pooling layer for 2D inputs.
See Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition,
K. He, X. Zhang, S. Ren, J. Sun
# Arguments
    pool_size: int
        Size of pooling region to use. pool_size = 7 will result in a 7x7 region.
    num_rois: number of regions of interest to be used
# Input shape
    list of two 4D tensors [X_img,X_roi] with shape:
    X_img:
    `(1, rows, cols, channels)`
    X_roi:
    `(1,num_rois,4)` list of rois, with ordering (x,y,w,h)
# Output shape
    3D tensor with shape:
    `(1, num_rois, channels, pool_size, pool_size)`
'''
def __init__(self, pool_size, num_rois, **kwargs):

    self.dim_ordering = K.image_dim_ordering()
    self.pool_size = pool_size
    self.num_rois = num_rois

    super(RoiPoolingConv, self).__init__(**kwargs)

def build(self, input_shape):
    self.nb_channels = input_shape[0][3]   

def compute_output_shape(self, input_shape):
    return None, self.num_rois, self.pool_size, self.pool_size, self.nb_channels

def call(self, x, mask=None):

    assert(len(x) == 2)

    # x[0] is image with shape (rows, cols, channels)
    img = x[0]

    # x[1] is roi with shape (num_rois,4) with ordering (x,y,w,h)
    rois = x[1]

    input_shape = K.shape(img)

    outputs = []

    for roi_idx in range(self.num_rois):

        x = rois[0, roi_idx, 0]
        y = rois[0, roi_idx, 1]
        w = rois[0, roi_idx, 2]
        h = rois[0, roi_idx, 3]

        x = K.cast(x, 'int32')
        y = K.cast(y, 'int32')
        w = K.cast(w, 'int32')
        h = K.cast(h, 'int32')

        # Resized roi of the image to pooling size (7x7)
        rs = tf.image.resize_images(img[:, y:y+h, x:x+w, :], (self.pool_size, self.pool_size))
        outputs.append(rs)


    final_output = K.concatenate(outputs, axis=0)

    # Reshape to (1, num_rois, pool_size, pool_size, nb_channels)
    # Might be (1, 4, 7, 7, 3)
    final_output = K.reshape(final_output, (1, self.num_rois, self.pool_size, self.pool_size, self.nb_channels))

    # permute_dimensions is similar to transpose
    final_output = K.permute_dimensions(final_output, (0, 1, 2, 3, 4))

    return final_output


def get_config(self):
    config = {'pool_size': self.pool_size,
              'num_rois': self.num_rois}
    base_config = super(RoiPoolingConv, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

I have looked at most of the resources out there and almost all of them suggest to comment out this layer. But since this layer is important for object detection, I was wondering if a workaround is possible or not.

The complete traceback of error (note: I've saved filename as freezekeras.py , contents are same as keras_to_tensorflow.py ):

Using TensorFlow backend.
Traceback (most recent call last):


File "freezekeras.py", line 181, in <module>
    app.run(main)
  File "/usr/local/lib/python3.5/dist-packages/absl/app.py", line 300, in run
    _run_main(main, args)
  File "/usr/local/lib/python3.5/dist-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "freezekeras.py", line 127, in main
    model = load_model(FLAGS.input_model, FLAGS.input_model_json, FLAGS.input_model_yaml)
  File "freezekeras.py", line 105, in load_model
    raise wrong_file_err
  File "freezekeras.py", line 62, in load_model
    model = keras.models.load_model(input_model_path)
  File "/usr/local/lib/python3.5/dist-packages/keras/engine/saving.py", line 419, in load_model
    model = _deserialize_model(f, custom_objects, compile)
  File "/usr/local/lib/python3.5/dist-packages/keras/engine/saving.py", line 225, in _deserialize_model
    model = model_from_config(model_config, custom_objects=custom_objects)
  File "/usr/local/lib/python3.5/dist-packages/keras/engine/saving.py", line 458, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "/usr/local/lib/python3.5/dist-packages/keras/layers/__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "/usr/local/lib/python3.5/dist-packages/keras/utils/generic_utils.py", line 145, in deserialize_keras_object
    list(custom_objects.items())))
  File "/usr/local/lib/python3.5/dist-packages/keras/engine/network.py", line 1022, in from_config
    process_layer(layer_data)
  File "/usr/local/lib/python3.5/dist-packages/keras/engine/network.py", line 1008, in process_layer
    custom_objects=custom_objects)
  File "/usr/local/lib/python3.5/dist-packages/keras/layers/__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "/usr/local/lib/python3.5/dist-packages/keras/utils/generic_utils.py", line 138, in deserialize_keras_object
    ': ' + class_name)
ValueError: Unknown layer: RoiPoolingConv

Answer 1

Try to specify the custom layer explicitly:

model = load_model('my_model.h5', custom_objects={'RoiPoolingConv': RoiPoolingConv})

Obviously, you have to re-write the keras_to_tensorflow.py script. See Handling custom layers (or other custom objects) in saved models section under Keras FAQ .

Answer 2

Solution

1. specify custom layer while loading model in keras_to_tensorflow.py

    model = keras.models.load_model(input_model_path, custom_objects={'RoiPoolingConv':RoiPoolingConv})

2. import RoiPoolingConv.py to keras_to_tensorflow project
3. specify default pool_size, num_rois for RoiPoolingConv

    def __init__(self, pool_size = 7, num_rois = 32, **kwargs):