ConvNeXt models missing gradients when compiling with `jit_compile=True`

[ianstenbit]

Example error:

WARNING:tensorflow:Gradients do not exist for variables 
['convnext_xlarge_stage_0_block_0_depthwise_conv/kernel:0', 'convnext_xlarge_stage_0_block_1_depthwise_conv/kernel:0', 'convnext_xlarge_stage_0_block_2_depthwise_conv/kernel:0', 'convnext_xlarge_stage_1_block_0_depthwise_conv/kernel:0', 'convnext_xlarge_stage_1_block_1_depthwise_conv/kernel:0', 'convnext_xlarge_stage_1_block_2_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_0_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_1_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_2_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_3_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_4_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_5_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_6_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_7_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_8_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_9_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_10_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_11_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_12_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_13_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_14_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_15_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_16_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_17_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_18_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_19_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_20_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_21_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_22_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_23_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_24_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_25_depthwise_conv/kernel:0', 'convnext_xlarge_stage_2_block_26_depthwise_conv/kernel:0', 'convnext_xlarge_stage_3_block_0_depthwise_conv/kernel:0', 'convnext_xlarge_stage_3_block_1_depthwise_conv/kernel:0', 'convnext_xlarge_stage_3_block_2_depthwise_conv/kernel:0'] 
when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss`argument?

[ianstenbit]

This is occurring when trying to train on TPU.

With jit_compile=False, training still fails, but with a different error:

tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc:113] Reshape's input dynamic dimension is decomposed into multiple output dynamic dimensions, but the constraint is ambiguous and XLA can't infer the output dimension %reshape.2192 = f32[64,56,56,1024]{3,2,1,0} reshape(f32[200704,1024]{1,0} %transpose.2191), metadata={op_type="Reshape" op_name="convnext_xlarge/convnext_xlarge_stage_0_block_0_pointwise_conv_1/Tensordot"}.
2022-09-02 17:21:54.799608: F tensorflow/core/tpu/kernels/tpu_program_group.cc:86] Check failed: xla_tpu_programs.size() > 0 (0 vs. 0)

[bhack]

Is it the same with drop_path_rate=False?

[ianstenbit]

@bhack yes, same error.

According to the error message above, it looks like the op with the shape error is convnext_xlarge/convnext_xlarge_stage_0_block_0_pointwise_conv_1/Tensordot, which is before the ConvNeXtBlocks where drop_path_rate is used.

[LukeWood]

@sayakpaul

[bhack]

Are you using

keras-cv/keras_cv/datasets/imagenet/load.py

Line 111 in b966b18

dataset = dataset.batch(batch_size)

for the training?

In that case have you tried to activate the drop_remainder=True param?

https://www.tensorflow.org/api_docs/python/tf/data/Dataset#batch

[sayakpaul]

Happens for all the variants or only XLarge?

[sayakpaul]

Need some reproducible code to debug it.

[ianstenbit]

@sayakpaul this is happening for all 3 that I've tested (Base, Large, XLarge)

[ianstenbit]

Need some reproducible code to debug it.

I am running this using our ImageNet classification training script (with this TPU training flag)

Invocation:

python3 -i ../keras-cv/examples/training/classification/imagenet/basic_training.py -imagenet_path=[redacted] --backup_path=[redacted] --weights_path=[redacted] --tensorboard_path=[redacted] --model_name=ConvNeXtXLarge --batch_size=512 --use_xla=False --use_tpu=True --model_kwargs="{'layer_scale_init_value': 1e-6, 'drop_path_rate': False}"

[sayakpaul]

A Colab Notebook is the only way for me to reproduce this. Could you provide one?

[ianstenbit]

A Colab Notebook is the only way for me to reproduce this. Could you provide one?

I'm unable to get a Colab notebook set up running KerasCV at this time, as KerasCV requires TF 2.9, and Colab is currently running TF 2.8.

I will share a Colab notebook with this example once Colab is updated to a sufficient version of TF

[bhack]

@ianstenbit Ok, in the meantime, as you have the setup ready, can you test it with drop_remainder=True?

[ianstenbit]

@ianstenbit Ok, in the meantime, as you have the setup ready, can you test it with drop_remainder=True?

Sorry I forgot to reply to this earlier -- I gave that a shot and it did not fix the problem.

[bhack]

Sorry I forgot to reply to this earlier -- I gave that a shot and it did not fix the problem.

Thanks.

[sayakpaul]

A Colab Notebook is the only way for me to reproduce this. Could you provide one?

I'm unable to get a Colab notebook set up running KerasCV at this time, as KerasCV requires TF 2.9, and Colab is currently running TF 2.8.

I will share a Colab notebook with this example once Colab is updated to a sufficient version of TF

Appreciate it, thanks.