Failed to initialize CoAtNet0 class

[Iqrar99]

After installing keras-cv-attention-models library using pip. I tried to import CoAtNet.

from keras_cv_attention_models import coatnet

co_model = coatnet.CoAtNet0()

Then it gives me this error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_554883/3784669224.py in <module>
----> 1 co_model = coatnet.CoAtNet0()

/opt/conda/lib/python3.8/site-packages/keras_cv_attention_models/coatnet/coatnet.py in CoAtNet0(input_shape, num_classes, drop_connect_rate, classifier_activation, pretrained, **kwargs)
    213     out_channels = [96, 192, 384, 768]
    214     stem_width = 64
--> 215     return CoAtNet(**locals(), model_name="coatnet0", **kwargs)
    216 
    217 

/opt/conda/lib/python3.8/site-packages/keras_cv_attention_models/coatnet/coatnet.py in CoAtNet(num_blocks, out_channels, stem_width, block_types, strides, expansion, se_ratio, head_dimension, use_dw_strides, bn_act_first, input_shape, num_classes, activation, drop_connect_rate, classifier_activation, dropout, pretrained, model_name, kwargs)
    192                 )
    193             else:
--> 194                 nn = res_mhsa(nn, out_channel, conv_short_cut, stride, head_dimension, block_drop_rate, activation=activation, name=name)
    195                 nn = res_ffn(nn, expansion=expansion, drop_rate=block_drop_rate, activation=activation, name=name + "ffn_")
    196 

/opt/conda/lib/python3.8/site-packages/keras_cv_attention_models/coatnet/coatnet.py in res_mhsa(inputs, output_channel, conv_short_cut, strides, head_dimension, drop_rate, activation, name)
    138         nn = keras.layers.MaxPool2D(pool_size=2, strides=strides, padding="SAME", name=name + "pool")(nn)
    139     num_heads = nn.shape[-1] // head_dimension
--> 140     nn = mhsa_with_multi_head_relative_position_embedding(nn, num_heads=num_heads, key_dim=head_dimension, out_shape=output_channel, name=name + "mhsa_")
    141     nn = drop_block(nn, drop_rate=drop_rate, name=name)
    142     # print(f"{name = }, {inputs.shape = }, {shortcut.shape = }, {nn.shape = }")

/opt/conda/lib/python3.8/site-packages/keras_cv_attention_models/coatnet/coatnet.py in mhsa_with_multi_head_relative_position_embedding(inputs, num_heads, key_dim, global_query, out_shape, out_weight, qkv_bias, out_bias, attn_dropout, name)
     23     _, hh, ww, cc = inputs.shape
     24     key_dim = key_dim if key_dim > 0 else cc // num_heads
---> 25     qk_scale = float(1.0 / tf.math.sqrt(tf.cast(key_dim, "float32")))
     26     out_shape = cc if out_shape is None or not out_weight else out_shape
     27     qk_out = num_heads * key_dim

TypeError: float() argument must be a string or a number, not 'Tensor'

I used Tensorflow 2.10 and Python 3.8. Tried it using Tensorflow 2.9 but no luck. Is there a way to fix it?

[Iqrar99]

For now i can obtain the model via direct download using wget

!wget -O pretrained_coatnet224.h5 https://github.com/leondgarse/keras_cv_attention_models/releases/download/coatnet/coatnet0_224_imagenet.h5

Then load the model manually

coatnet_model_path = './pretrained_coatnet224.h5'
co_model = tf.keras.models.load_model(coatnet_model_path)

[leondgarse]

• It's rather weird... You may try if replacing coatnet.py L25 qk_scale = float(1.0 / tf.math.sqrt(tf.cast(key_dim, "float32"))) with qk_scale = float(1.0 / key_dim ** 0.5) works for you.
• This colab kecam_test.ipynb contains a CoAtNet0 basic tests there, and it works rather normal...
• For the version, I'm curently using Python 3.8.5 and TF 2.10.0, and in the colab it's python 3.8.16 + TF 2.9.2. It just doesn't that matters, other versions should also work fine.

[Iqrar99]

Noted with thanks. I'll try it later.