Keras Lambda Layer: can't pickle _thread.RLock objects

Question

I am trying to save the full keras model. I tried to replicate the following error in a smaller example but I failed, so here is the actual code causing the error in the full model.

The MaskRNN model for image segmentation has a lambda layer defined as:

gt_boxes = KL.Lambda(lambda x: norm_boxes_graph(x, K.shape(input_image)[1:3]))(input_gt_boxes)

This layer calls the method below:

def norm_boxes_graph(self, boxes, shape):
        """Converts boxes from pixel coordinates to normalized coordinates.
        boxes: [..., (y1, x1, y2, x2)] in pixel coordinates
        shape: [..., (height, width)] in pixels

        Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
        coordinates it's inside the box.

        Returns:
            [..., (y1, x1, y2, x2)] in normalized coordinates
        """
        h, w = tf.split(tf.cast(shape, tf.float32), 2)
        scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
        shift = tf.constant([0., 0., 1., 1.])
        fin = tf.divide(boxes - shift, scale)
        return fin

When trying to save the full model(not only the weights) it throughs the error:

can't pickle _thread.RLock objects

Checking for serializability:

Checking Serializability of {'class_name': 'Lambda', 
'config': {'name': 'norm_boxes_lambda', 'trainable': True, 'dtype': 'float32', 'function': ('4wEAAAAAAAAAAQAAAAQAAAATAAAAcw4AAACIAHwAiAGIAoMBgwJTACkBTqkAKQHaAXgpA9oPZ2V0\nX25vcm1fYm94ZXNf2hBnZXRfc2hhcGVfaW1hZ2Vf2gtpbnB1dF9pbWFnZXIBAAAA+lEvaG9tZS9p\nYWkvbWluaWNvbmRhMi9lbnZzL3B5dG9yY2gyL2xpYi9weXRob24zLjcvc2l0ZS1wYWNrYWdlcy9N\nYXNrX1JDTk4vbW9kZWwucHnaCDxsYW1iZGE+mwcAAPMAAAAA\n', None, (<function MaskRCNN.build.<locals>.get_norm_boxes_ at 0x7f272b8b59e0>, <function MaskRCNN.build.<locals>.get_shape_image_ at 0x7f270c4817a0>, <tf.Tensor 'input_image:0' shape=(None, None, None, 3) dtype=float32>)), 
'function_type': 'lambda', 
'module': 'Mask_RCNN.model', 
'output_shape': None, 
'output_shape_type': 'raw', 'output_shape_module': None, 
'arguments': {}}}


!!! FAIL serialization: can't pickle _thread.RLock objects
WARNING: Did not find non-serializable object in {'class_name': 'Lambda', 'config': {'name': 'norm_boxes_lambda', 'trainable': True, 'dtype': 'float32', 'function': ('4wEAAAAAAAAAAQAAAAQAAAATAAAAcw4AAACIAHwAiAGIAoMBgwJTACkBTqkAKQHaAXgpA9oPZ2V0\nX25vcm1fYm94ZXNf2hBnZXRfc2hhcGVfaW1hZ2Vf2gtpbnB1dF9pbWFnZXIBAAAA+lEvaG9tZS9p\nYWkvbWluaWNvbmRhMi9lbnZzL3B5dG9yY2gyL2xpYi9weXRob24zLjcvc2l0ZS1wYWNrYWdlcy9N\nYXNrX1JDTk4vbW9kZWwucHnaCDxsYW1iZGE+mwcAAPMAAAAA\n', None, (<function MaskRCNN.build.<locals>.get_norm_boxes_ at 0x7f272b8b59e0>, <function MaskRCNN.build.<locals>.get_shape_image_ at 0x7f270c4817a0>, <tf.Tensor 'input_image:0' shape=(None, None, None, 3) dtype=float32>)), 'function_type': 'lambda', 'module': 'Mask_RCNN.model', 'output_shape': None, 'output_shape_type': 'raw', 'output_shape_module': None, 'arguments': {}}}. This may be an oversight.

========================
Variable: 

    FailTuple(test [obj={'class_name': 'Lambda', 'config': {'name': 'norm_boxes_lambda', 'trainable': True, 'dtype': 'float32', 'function': ('4wEAAAAAAAAAAQAAAAQAAAATAAAAcw4AAACIAHwAiAGIAoMBgwJTACkBTqkAKQHaAXgpA9oPZ2V0\nX25vcm1fYm94ZXNf2hBnZXRfc2hhcGVfaW1hZ2Vf2gtpbnB1dF9pbWFnZXIBAAAA+lEvaG9tZS9p\nYWkvbWluaWNvbmRhMi9lbnZzL3B5dG9yY2gyL2xpYi9weXRob24zLjcvc2l0ZS1wYWNrYWdlcy9N\nYXNrX1JDTk4vbW9kZWwucHnaCDxsYW1iZGE+mwcAAPMAAAAA\n', None, (<function MaskRCNN.build.<locals>.get_norm_boxes_ at 0x7f272b8b59e0>, <function MaskRCNN.build.<locals>.get_shape_image_ at 0x7f270c4817a0>, <tf.Tensor 'input_image:0' shape=(None, None, None, 3) dtype=float32>)), 'function_type': 'lambda', 'module': 'Mask_RCNN.model', 'output_shape': None, 'output_shape_type': 'raw', 'output_shape_module': None, 'arguments': {}}}, parent=None])

was found to be non-serializable. There may be multiple other undetected variables that were non-serializable. 
Consider either removing the instantiation/imports of these variables or moving the instantiation into the scope of the function/class. 
If you have any suggestions on how to improve this error message, please reach out to the Ray developers on github.com/ray-project/ray/issues/

Useful related info: Keras Custom Layer gives errors when saving the full model

I also used unsuccessfully the closure method described here: Checkpointing keras model: TypeError: can't pickle _thread.lock objects

Any help is greatly appreciated!

Keras Lambda Layer: can't pickle _thread.RLock objects

0 Answers0