InvalidArgumentError when using RandomCrop on M1 Pro GPU

Hi,

I am using a Mac with M1 Pro. I want to use RandomCrop from tensorflow.keras.layers but while training I get the error below. If I understood correctly, it seems that RngReadAndSkip is not implemented for the GPU.

InvalidArgumentError: Cannot assign a device for operation model/data_augmentation/random_crop/cond/model_data_augmentation_random_crop_cond_input_1/_6: Could not satisfy explicit device specification '' because the node {{colocation_node model/data_augmentation/random_crop/cond/model_data_augmentation_random_crop_cond_input_1/_6}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:GPU:0]. 
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' assigned_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' resource_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[]
RngReadAndSkip: CPU 
Identity: GPU CPU 
Switch: GPU CPU 
_Arg: GPU CPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  model_data_augmentation_random_crop_cond_input_1 (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  model/data_augmentation/random_crop/cond/model_data_augmentation_random_crop_cond_input_1/_6 (Switch) 
  Func/model/data_augmentation/random_crop/cond/then/_0/input/_47 (Identity) 
  model/data_augmentation/random_crop/cond/then/_0/model/data_augmentation/random_crop/cond/stateful_uniform/RngReadAndSkip (RngReadAndSkip) 
  Func/model/data_augmentation/random_crop/cond/else/_1/input/_52 (Identity) 

Python version:

$ python --version --version
Python 3.8.13 (default, Mar 28 2022, 06:13:39) 
[Clang 12.0.0 ]

Libraries used:

$ conda list  | grep tensorflow
tensorflow-addons         0.17.0                   pypi_0    pypi
tensorflow-deps           2.9.0                         0    apple
tensorflow-estimator      2.9.0                    pypi_0    pypi
tensorflow-macos          2.9.2                    pypi_0    pypi
tensorflow-metal          0.5.0                    pypi_0    pypi

Is there any workaround? Or anything I can do to help fixing this?

Thanks

Answered by Frameworks Engineer in 724511022

RngReadAndSkip is registered on the GPU in tensorflow-metal==0.5.1.

Hi @EdoAbati

The issue comes up because the RngReadAndSkip is missing its GPU implementation at the moment. We are working on implementing it. In the meanwhile the workaround would be to separate the RandomCrop to run on the CPU using the block with tf.device('/cpu:0'):. We will update this thread soon once the fix is out.

Accepted Answer

RngReadAndSkip is registered on the GPU in tensorflow-metal==0.5.1.

InvalidArgumentError when using RandomCrop on M1 Pro GPU
 
 
Q