{
  "id": 140567,
  "title": "XLA compilation error with custom layer",
  "url": "/competitions/flower-classification-with-tpus/discussion/140567",
  "author_name": "",
  "post_date": "2020-04-02T12:45:27.685573600Z",
  "votes": null,
  "comment_count": 5,
  "views": 0,
  "content": "<p>I've written my own custom layer and I'm having an issue training it. When training begins, I get the following error:</p>\n\n<p><code>\nCompilation failure: XLA can't deduce compile time constant output shape for strided slice: [?,512,512,3], output shape must be a compile-time constant\n</code></p>\n\n<p>My custom layer looks something like this:</p>\n\n<p>```\nclass MyCustomLayer(Layer):</p>\n\n<pre><code>def __init__(self, name=None):\n    super(MyCustomLayer, self).__init__(name=name)\n\n@tf.function\ndef call(self, inputs):\n    \"\"\"\n    inputs has two tensors : [images, image_dependent_params]\n    \"\"\"\n    images, image_dependent_params = inputs\n\n    outputs = tf.zeros_like(images)\n    for i in tf.range(tf.shape(images)[0]):\n        # compute outputs[i] based on images[i] and image_dependent_params[i]\n        # set outputs[i] by doing zero padded concatenation\n\n    return outputs\n</code></pre>\n\n<p>```</p>\n\n<p>Some points about my custom layer:</p>\n\n<ul>\n<li><p>I do zero-padded concatenation because I couldn't figure out how to do slice assignment on the <code>outputs</code> tensor. Zero-padded because I can't change the shape of <code>outputs</code> over loop iterations. See <a href=\"https://www.tensorflow.org/tutorials/customization/performance#consistent_shapes_and_types\">here</a> if you are still curious about this.</p></li>\n<li><p>I use @tf.function because Tensorflow didn't seem to be happy about the Python flow control in the for loop.</p></li>\n<li><p>To be honest, I'd love to avoid the loop altogether, but I don't think that's possible because of <code>image_dependent_params</code> and the operations I am doing.</p></li>\n</ul>",
  "messages": [
    {
      "id": "795125",
      "postDate": "04/02/2020 12:45:27",
      "content": "<p>I've written my own custom layer and I'm having an issue training it. When training begins, I get the following error:</p>\n\n<p><code>\nCompilation failure: XLA can't deduce compile time constant output shape for strided slice: [?,512,512,3], output shape must be a compile-time constant\n</code></p>\n\n<p>My custom layer looks something like this:</p>\n\n<p>```\nclass MyCustomLayer(Layer):</p>\n\n<pre><code>def __init__(self, name=None):\n    super(MyCustomLayer, self).__init__(name=name)\n\n@tf.function\ndef call(self, inputs):\n    \"\"\"\n    inputs has two tensors : [images, image_dependent_params]\n    \"\"\"\n    images, image_dependent_params = inputs\n\n    outputs = tf.zeros_like(images)\n    for i in tf.range(tf.shape(images)[0]):\n        # compute outputs[i] based on images[i] and image_dependent_params[i]\n        # set outputs[i] by doing zero padded concatenation\n\n    return outputs\n</code></pre>\n\n<p>```</p>\n\n<p>Some points about my custom layer:</p>\n\n<ul>\n<li><p>I do zero-padded concatenation because I couldn't figure out how to do slice assignment on the <code>outputs</code> tensor. Zero-padded because I can't change the shape of <code>outputs</code> over loop iterations. See <a href=\"https://www.tensorflow.org/tutorials/customization/performance#consistent_shapes_and_types\">here</a> if you are still curious about this.</p></li>\n<li><p>I use @tf.function because Tensorflow didn't seem to be happy about the Python flow control in the for loop.</p></li>\n<li><p>To be honest, I'd love to avoid the loop altogether, but I don't think that's possible because of <code>image_dependent_params</code> and the operations I am doing.</p></li>\n</ul>",
      "rawMarkdown": "I've written my own custom layer and I'm having an issue training it. When training begins, I get the following error:\n\n```\nCompilation failure: XLA can't deduce compile time constant output shape for strided slice: [?,512,512,3], output shape must be a compile-time constant\n```\n\nMy custom layer looks something like this:\n\n```\nclass MyCustomLayer(Layer):\n\n    def __init__(self, name=None):\n        super(MyCustomLayer, self).__init__(name=name)\n\n    @tf.function\n    def call(self, inputs):\n        \"\"\"\n        inputs has two tensors : [images, image_dependent_params]\n        \"\"\"\n        images, image_dependent_params = inputs\n\n        outputs = tf.zeros_like(images)\n        for i in tf.range(tf.shape(images)[0]):\n            # compute outputs[i] based on images[i] and image_dependent_params[i]\n            # set outputs[i] by doing zero padded concatenation\n\n        return outputs\n```\n\nSome points about my custom layer:\n\n- I do zero-padded concatenation because I couldn't figure out how to do slice assignment on the `outputs` tensor. Zero-padded because I can't change the shape of `outputs` over loop iterations. See [here](https://www.tensorflow.org/tutorials/customization/performance#consistent_shapes_and_types) if you are still curious about this.\n\n- I use @tf.function because Tensorflow didn't seem to be happy about the Python flow control in the for loop.\n\n- To be honest, I'd love to avoid the loop altogether, but I don't think that's possible because of `image_dependent_params` and the operations I am doing.",
      "votes": null
    },
    {
      "id": "795309",
      "postDate": "04/02/2020 15:55:59",
      "content": "<p>if you are already doing the work of padding and outputting a fixed shape, maybe Tensorflow has a hard time realizing that the end shape is indeed fixed. Have you tried forcing it with a tf.reshape ?</p>",
      "rawMarkdown": "if you are already doing the work of padding and outputting a fixed shape, maybe Tensorflow has a hard time realizing that the end shape is indeed fixed. Have you tried forcing it with a tf.reshape ?",
      "votes": null
    },
    {
      "id": "795386",
      "postDate": "04/02/2020 17:05:08",
      "content": "<p>Thanks for the response. I tried that just then and unfortunately the same error comes up.</p>\n\n<p>It might also help to know that just running model.predict does work (at least it does in my local machine with CPU - will try it on the TPU shortly). \nIt's just model.fit which falls apart.</p>",
      "rawMarkdown": "Thanks for the response. I tried that just then and unfortunately the same error comes up.\n\nIt might also help to know that just running model.predict does work (at least it does in my local machine with CPU - will try it on the TPU shortly). \nIt's just model.fit which falls apart.",
      "votes": null
    },
    {
      "id": "795429",
      "postDate": "04/02/2020 17:54:24",
      "content": "<p>Ah I think it's to do with TPU.</p>\n\n<p>Okay, so a correction on the above. <code>model.predict</code> does not work on TPU either.\nI can verify it's certainly a TPU/CPU thing because if I pull my model definition out of <code>with strategy.scope()</code>, I'm able to run <code>model.predict</code>.</p>\n\n<p>That's a shame. The whole reason I wrote this funny loop was because <code>tf.images.crop_and_resize</code> is not supported on TPU yet.</p>",
      "rawMarkdown": "Ah I think it's to do with TPU.\n\nOkay, so a correction on the above. `model.predict` does not work on TPU either.\nI can verify it's certainly a TPU/CPU thing because if I pull my model definition out of `with strategy.scope()`, I'm able to run `model.predict`.\n\nThat's a shame. The whole reason I wrote this funny loop was because `tf.images.crop_and_resize` is not supported on TPU yet.",
      "votes": null
    },
    {
      "id": "795627",
      "postDate": "04/02/2020 22:15:00",
      "content": "<p>Got it. If your goal is to crop and resize images in funny ways you can:\n- pre-process the dataset into a new one\n- if you want to do this this in real time, do it in the tf.data.Dataset code instead of a Layer. Dataset runs on the CPU side of the TPU and has fewer limitations.</p>\n\n<p>I also think there must be a way to write this with fixed outputs.</p>",
      "rawMarkdown": "Got it. If your goal is to crop and resize images in funny ways you can:\n- pre-process the dataset into a new one\n- if you want to do this this in real time, do it in the tf.data.Dataset code instead of a Layer. Dataset runs on the CPU side of the TPU and has fewer limitations.\n\nI also think there must be a way to write this with fixed outputs.",
      "votes": null
    },
    {
      "id": "795638",
      "postDate": "04/02/2020 22:29:22",
      "content": "<p>Well actually I need to do the cropping in a layer internal to the model. I noticed a comment on this <a href=\"https://github.com/tensorflow/tensorflow/issues/23337\">feature request</a> claiming that it wasn't really possible to do the <code>crop_and_resize</code> op with the XLA compiler, and I'm guessing the reasons for it probably translate to my attempted workaround. I'm now looking into the suggestion to use <code>tpu.outside_compilation</code> although this is leading me down a whole other rabbit hole of errors.</p>",
      "rawMarkdown": "Well actually I need to do the cropping in a layer internal to the model. I noticed a comment on this [feature request](https://github.com/tensorflow/tensorflow/issues/23337) claiming that it wasn't really possible to do the `crop_and_resize` op with the XLA compiler, and I'm guessing the reasons for it probably translate to my attempted workaround. I'm now looking into the suggestion to use `tpu.outside_compilation` although this is leading me down a whole other rabbit hole of errors.",
      "votes": null
    }
  ],
  "comments": [
    {
      "id": 795309,
      "author_name": "mgorner",
      "author_url": "",
      "post_date": "04/02/2020 15:55:59",
      "content": "<p>if you are already doing the work of padding and outputting a fixed shape, maybe Tensorflow has a hard time realizing that the end shape is indeed fixed. Have you tried forcing it with a tf.reshape ?</p>",
      "votes": null,
      "replies": [
        {
          "id": 795386,
          "author_name": "alexandersoare",
          "author_url": "",
          "post_date": "04/02/2020 17:05:08",
          "content": "<p>Thanks for the response. I tried that just then and unfortunately the same error comes up.</p>\n\n<p>It might also help to know that just running model.predict does work (at least it does in my local machine with CPU - will try it on the TPU shortly). \nIt's just model.fit which falls apart.</p>",
          "votes": null,
          "replies": []
        },
        {
          "id": 795429,
          "author_name": "alexandersoare",
          "author_url": "",
          "post_date": "04/02/2020 17:54:24",
          "content": "<p>Ah I think it's to do with TPU.</p>\n\n<p>Okay, so a correction on the above. <code>model.predict</code> does not work on TPU either.\nI can verify it's certainly a TPU/CPU thing because if I pull my model definition out of <code>with strategy.scope()</code>, I'm able to run <code>model.predict</code>.</p>\n\n<p>That's a shame. The whole reason I wrote this funny loop was because <code>tf.images.crop_and_resize</code> is not supported on TPU yet.</p>",
          "votes": null,
          "replies": []
        },
        {
          "id": 795627,
          "author_name": "mgorner",
          "author_url": "",
          "post_date": "04/02/2020 22:15:00",
          "content": "<p>Got it. If your goal is to crop and resize images in funny ways you can:\n- pre-process the dataset into a new one\n- if you want to do this this in real time, do it in the tf.data.Dataset code instead of a Layer. Dataset runs on the CPU side of the TPU and has fewer limitations.</p>\n\n<p>I also think there must be a way to write this with fixed outputs.</p>",
          "votes": null,
          "replies": []
        },
        {
          "id": 795638,
          "author_name": "alexandersoare",
          "author_url": "",
          "post_date": "04/02/2020 22:29:22",
          "content": "<p>Well actually I need to do the cropping in a layer internal to the model. I noticed a comment on this <a href=\"https://github.com/tensorflow/tensorflow/issues/23337\">feature request</a> claiming that it wasn't really possible to do the <code>crop_and_resize</code> op with the XLA compiler, and I'm guessing the reasons for it probably translate to my attempted workaround. I'm now looking into the suggestion to use <code>tpu.outside_compilation</code> although this is leading me down a whole other rabbit hole of errors.</p>",
          "votes": null,
          "replies": []
        }
      ]
    }
  ],
  "raw_markdown_by_id": {
    "795125": "I've written my own custom layer and I'm having an issue training it. When training begins, I get the following error:\n\n```\nCompilation failure: XLA can't deduce compile time constant output shape for strided slice: [?,512,512,3], output shape must be a compile-time constant\n```\n\nMy custom layer looks something like this:\n\n```\nclass MyCustomLayer(Layer):\n\n    def __init__(self, name=None):\n        super(MyCustomLayer, self).__init__(name=name)\n\n    @tf.function\n    def call(self, inputs):\n        \"\"\"\n        inputs has two tensors : [images, image_dependent_params]\n        \"\"\"\n        images, image_dependent_params = inputs\n\n        outputs = tf.zeros_like(images)\n        for i in tf.range(tf.shape(images)[0]):\n            # compute outputs[i] based on images[i] and image_dependent_params[i]\n            # set outputs[i] by doing zero padded concatenation\n\n        return outputs\n```\n\nSome points about my custom layer:\n\n- I do zero-padded concatenation because I couldn't figure out how to do slice assignment on the `outputs` tensor. Zero-padded because I can't change the shape of `outputs` over loop iterations. See [here](https://www.tensorflow.org/tutorials/customization/performance#consistent_shapes_and_types) if you are still curious about this.\n\n- I use @tf.function because Tensorflow didn't seem to be happy about the Python flow control in the for loop.\n\n- To be honest, I'd love to avoid the loop altogether, but I don't think that's possible because of `image_dependent_params` and the operations I am doing.",
    "795309": "if you are already doing the work of padding and outputting a fixed shape, maybe Tensorflow has a hard time realizing that the end shape is indeed fixed. Have you tried forcing it with a tf.reshape ?",
    "795386": "Thanks for the response. I tried that just then and unfortunately the same error comes up.\n\nIt might also help to know that just running model.predict does work (at least it does in my local machine with CPU - will try it on the TPU shortly). \nIt's just model.fit which falls apart.",
    "795429": "Ah I think it's to do with TPU.\n\nOkay, so a correction on the above. `model.predict` does not work on TPU either.\nI can verify it's certainly a TPU/CPU thing because if I pull my model definition out of `with strategy.scope()`, I'm able to run `model.predict`.\n\nThat's a shame. The whole reason I wrote this funny loop was because `tf.images.crop_and_resize` is not supported on TPU yet.",
    "795627": "Got it. If your goal is to crop and resize images in funny ways you can:\n- pre-process the dataset into a new one\n- if you want to do this this in real time, do it in the tf.data.Dataset code instead of a Layer. Dataset runs on the CPU side of the TPU and has fewer limitations.\n\nI also think there must be a way to write this with fixed outputs.",
    "795638": "Well actually I need to do the cropping in a layer internal to the model. I noticed a comment on this [feature request](https://github.com/tensorflow/tensorflow/issues/23337) claiming that it wasn't really possible to do the `crop_and_resize` op with the XLA compiler, and I'm guessing the reasons for it probably translate to my attempted workaround. I'm now looking into the suggestion to use `tpu.outside_compilation` although this is leading me down a whole other rabbit hole of errors."
  },
  "source": "meta"
}