python - TPU 上自定义 Keras 层的 XLA 编译错误

Question

我已经编写了自己的自定义层，但在训练它时遇到了问题。训练开始时，我收到以下错误：

编译失败：XLA 无法推断跨步切片的编译时间常数输出形状：[?,512,512,3]，输出形状必须是编译时间常数

我的自定义图层看起来像这样（将在后期脚本中提供完整的图层）：

class MyCustomLayer(Layer):

    def __init__(self, name=None):
        super(MyCustomLayer, self).__init__(name=name)

    @tf.function
    def call(self, inputs):
        """
        inputs has two tensors : [images, image_dependent_params]
        """
        images, image_dependent_params = inputs

        outputs = tf.zeros_like(images)
        for i in tf.range(tf.shape(images)[0]):
            # compute outputs[i] based on images[i] and image_dependent_params[i]
            # set outputs[i] by doing zero padded concatenation

        return outputs

该循环是针对 TPU 不支持的crop_and_resize操作的解决方法。

PS：如果你有兴趣，这里是它的完整代码：

class AttentionLocalization(Layer):

    def __init__(self, name=None):
        super(AttentionLocalization, self).__init__(name=name)

    @tf.function
    def call(self, inputs):
        """
        inputs has two tensors : [images, rois]
        """
        images, rois = inputs
        tx, ty, tl = rois[:,0], rois[:,1], rois[:,2]
        tl = tf.clip_by_value(tl, 1/3, 1)
        tx = tf.clip_by_value(tx, 0, 1 - tl)
        ty = tf.clip_by_value(ty, 0, 1 - tl)

        size = tf.cast(tf.shape(images)[1], tf.float32)

        tx, ty, tl = rois[:,0], rois[:,1], rois[:,2]
        tl = tf.clip_by_value(tl, 1/3, 1)
        tx = tf.clip_by_value(tx, 0, 1 - tl)
        ty = tf.clip_by_value(ty, 0, 1 - tl)

        tl = tf.cast(tl * size, tf.int32)
        tx = tf.cast(tx * size, tf.int32)
        ty = tf.cast(ty * size, tf.int32)

        outputs = tf.zeros_like(images)
        batch_size = tf.shape(images)[0]
        for i in tf.range(batch_size):
            resized = tf.image.resize(images[i][tx[i]:tx[i]+tl[i], ty[i]:ty[i]+tl[i], :], tf.shape(images)[1:3])
            padding = tf.zeros([batch_size - 1 - i, tf.shape(images)[1], tf.shape(images)[2], tf.shape(images)[3]])
            outputs = tf.concat([outputs[:i], [resized], padding], axis=0)

        return outputs

编辑

我意识到我的自定义层确实可以在 CPU 上工作，但不能在我尝试训练模型的 TPU 上工作。

标签： pythontensorflowkeras