tensorflow / tpu

Reference models and tools for Cloud TPUs.
https://cloud.google.com/tpu/
Apache License 2.0
5.21k stars 1.77k forks source link

Detected unsupported operations when trying to compile graph __inference_one_step_on_data_4740[] on XLA_GPU_JIT: CropAndResize #1062

Open mattwilk25 opened 3 months ago

mattwilk25 commented 3 months ago

I'm getting the following error on cloud TPU when trying to use tf.image.crop_and_resize: Detected at node functional_1_1/crop_layer_1_1/CropAndResize defined at (most recent call last):

Detected at node functional_1_1/crop_layer_1_1/CropAndResize defined at (most recent call last): Detected unsupported operations when trying to compile graph __inference_one_step_on_data_4740[] on XLA_GPU_JIT: CropAndResize (No registered 'CropAndResize' OpKernel for XLA_GPU_JIT devices compatible with node {{node functional_1_1/crop_layer_1_1/CropAndResize}}){{node functional_1_1/crop_layer_1_1/CropAndResize}} The op is created at: File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main File "/usr/lib/python3.10/runpy.py", line 86, in _run_code File "/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py", line 37, in File "/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py", line 992, in launch_instance File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py", line 619, in start File "/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py", line 195, in start File "/usr/lib/python3.10/asyncio/base_events.py", line 603, in run_forever File "/usr/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once File "/usr/lib/python3.10/asyncio/events.py", line 80, in _run File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 685, in File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 738, in _run_callback File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 825, in inner File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 786, in run File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 361, in process_one File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 261, in dispatch_shell File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 539, in execute_request File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper File "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py", line 302, in do_execute File "/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py", line 539, in run_cell File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell File "/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code File "", line 6, in File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 117, in error_handler File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/tensorflow/trainer.py", line 318, in fit File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/tensorflow/trainer.py", line 121, in one_step_on_iterator File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/tensorflow/trainer.py", line 108, in one_step_on_data File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/tensorflow/trainer.py", line 51, in train_step File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 117, in error_handler File "/usr/local/lib/python3.10/dist-packages/keras/src/layers/layer.py", line 882, in __call__ File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 117, in error_handler File "/usr/local/lib/python3.10/dist-packages/keras/src/ops/operation.py", line 46, in __call__ File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 156, in error_handler File "/usr/local/lib/python3.10/dist-packages/keras/src/models/functional.py", line 175, in call File "/usr/local/lib/python3.10/dist-packages/keras/src/ops/function.py", line 171, in _run_through_graph File "/usr/local/lib/python3.10/dist-packages/keras/src/models/functional.py", line 556, in call File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 117, in error_handler File "/usr/local/lib/python3.10/dist-packages/keras/src/layers/layer.py", line 882, in __call__ File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 117, in error_handler File "/usr/local/lib/python3.10/dist-packages/keras/src/ops/operation.py", line 46, in __call__ File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 156, in error_handler File "", line 17, in call tf2xla conversion failed while converting __inference_one_step_on_data_4740[]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions. [[StatefulPartitionedCall]] [Op:__inference_one_step_on_iterator_4903] code context: import tensorflow as tf from tensorflow.keras import layers, Model class CropLayer(layers.Layer): def __init__(self, **kwargs): super(CropLayer, self).__init__(**kwargs) def call(self, inputs): image, crop_boxes = inputs batch_size = tf.shape(image)[0] num_boxes = tf.shape(crop_boxes)[1] crop_boxes = tf.reshape(crop_boxes, (batch_size * num_boxes, 4)) # Flatten the crop_boxes tensor box_indices = tf.range(batch_size) box_indices = tf.repeat(box_indices, repeats=num_boxes) img_size = tf.constant([48, 48], dtype=tf.int32) cropped_images = tf.image.crop_and_resize( image, crop_boxes, box_indices, img_size, method='bilinear', extrapolation_value=0.0, ) return cropped_images