在Mask RCNN模型中，我用一个自定义层替换了下面的Lambda层。当模型编译时，它不会在GPU上训练。它似乎在分配Workers之前停止在Epoch 1。我不确定我做错了什么。Lambda层不允许我保存模型，所以我必须使用另一种方法，但它不起作用。感谢任何帮助。

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

按自定义图层

gt_boxes = GtBoxesLayer(name='lambda_get_norm_boxes')([input_image, input_gt_boxes])

层代码为：

class GtBoxesLayer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(GtBoxesLayer, self).__init__(**kwargs)

    def call(self, input):
        return norm_boxes_graph(input[1], get_shape_image(input[0]))

    def get_config(self):
        config = super(GtBoxesLayer, self).get_config()
        return config
    
    @classmethod
    def from_config(cls, config):
        return cls(**config)


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

def get_shape_image_(input_image_):
     shape_= tf.shape(input_image_)
     return shape_[1:3]