Unable to run main example

I've tried to reproduce the experiment in my local environment, but even main.py will not execute to completion. The crashing occurs either when reading the ResNet56 model or when importing the keras classes.

Can you please specify which tensorlfow and keras version must be used for the system to be reproduceable?

Bellow I leave the output from my attempts on running main.py with two tensorflow versions (v2.16.2 and v2.15.1):

TF v2.16.2

python --version
>>> Python 3.12.4
python -c "import tensorflow as tf; print(tf.__version__)"
>>>  2.16.2
python main.py
>>> Traceback (most recent call last):
>>>   File "/root/personal/CKA-Layer-Pruning/main.py", line 192, in <module>
>>>     model = load_model('ResNet56')
>>>             ^^^^^^^^^^^^^^^^^^^^^^
>>>   File "/root/personal/CKA-Layer-Pruning/main.py", line 77, in load_model
>>>     from keras.utils.generic_utils import CustomObjectScope
>>> ModuleNotFoundError: No module named 'keras.utils.generic_utils'

After changing from keras.utils.generic_utils import CustomObjectScope to from keras.utils import CustomObjectScope (which should be allowed, to match docs of v2.16.1 on CustomObjectScope), the output becomes:

>>> Traceback (most recent call last):
>>>   File "/root/personal/CKA-Layer-Pruning/main.py", line 192, in <module>
>>>     model = load_model('ResNet56')
>>>             ^^^^^^^^^^^^^^^^^^^^^^
>>>   File "/root/personal/CKA-Layer-Pruning/main.py", line 94, in load_model
>>>     model = keras.models.model_from_json(f.read())
>>>             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>>>   File "/root/miniconda3/envs/tf2162/lib/python3.12/site-packages/keras/src/models/model.py", line 579, in model_from_json
>>>     return serialization_lib.deserialize_keras_object(
>>>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>>>   File "/root/miniconda3/envs/tf2162/lib/python3.12/site-packages/keras/src/saving/serialization_lib.py", line 694, in deserialize_keras_object
>>>     cls = _retrieve_class_or_fn(
>>>           ^^^^^^^^^^^^^^^^^^^^^^
>>>   File "/root/miniconda3/envs/tf2162/lib/python3.12/site-packages/keras/src/saving/serialization_lib.py", line 812, in _retrieve_class_or_fn
>>>     raise TypeError(
>>> TypeError: Could not locate class 'Model'. Make sure custom classes are decorated with `@keras.saving.register_keras_serializable()`. Full object config: {'keras_version': '2.1.2', 'config': {'input_layers': [['input_1', 0, 0]], 'layers': [{'inbound_nodes': [], 'config': {'batch_input_shape': [None, 32, 32, 3], 'sparse': False, 'dtype': 'float32', 'name': 'input_1'}, 'name': 'input_1', 'class_name': 'InputLayer'}, {'inbound_nodes': [[['input_1', 0, 0, {}]]], 'config': {'kernel_regularizer': {'config': {'l2': 9.999999747378752e-05, 'l1': 0.0}, 'class_name': 'L1L2'}, 'trainable': True, 'filters': 16, 'bias_constraint': None, 'activity_regularizer': None, 'use_bias': True, 'strides': [1, 1], 'name': 'conv2d_1', 'activation': 'linear', 'bias_initializer': {'config': {}, 'class_name': 'Zeros'}, 'padding': 'same', 'kernel_constraint': None, 'data_format': 'channels_last', 'dilation_rate': [1, 1], 'kernel_initializer': {'config': {'scale': 2.0, 'mode': 'fan_in', 'distribution': 'normal', 'seed': None}, 'class_name': 'VarianceScaling'}, 'kernel_size': [3, 3], 'bias_regularizer': None}, 'name': 'conv2d_1', 'class_name': 'Conv2D'}, {'inbound_nodes': [[['conv2d_1', 0, 0, {}]]], 'config': {'trainable': True, 'axis': -1, (...) 'bias_initializer': {'config': {}, 'class_name': 'Zeros'}, 'units': 10, 'kernel_constraint': None, 'bias_constraint': None, 'kernel_initializer': {'config': {'scale': 2.0, 'mode': 'fan_in', 'distribution': 'normal', 'seed': None}, 'class_name': 'VarianceScaling'}, 'bias_regularizer': None}, 'name': 'dense_1', 'class_name': 'Dense'}], 'output_layers': [['dense_1', 0, 0]], 'name': 'model_1'}, 'backend': 'tensorflow', 'class_name': 'Model'}

TF v2.15.1

python --version
>>> Python 3.9.13
python -c "import tensorflow as tf; print(tf.__version__)"
>>>  2.15.1
python -c "import keras; print(keras.__version__)"
>>> '2.15.0'
python main.py
>>>Traceback (most recent call last):
>>>  File "/root/personal/CKA-Layer-Pruning/main.py", line 192, in <module>
>>>    model = load_model('ResNet56')
>>>  File "/root/personal/CKA-Layer-Pruning/main.py", line 77, in load_model
>>>    from keras.utils.generic_utils import CustomObjectScope
>>> ModuleNotFoundError: No module named 'keras.utils.generic_utils'

IanPons / CKA-Layer-Pruning

Unable to run main example #1