from tensorflow.keras.layers import Conv2D, MaxPooling2D, GlobalAveragePooling2D, Dense, Dropout, LSTM, Reshape, Resizing, BatchNormalization, Flatten, TimeDistributed, Input, Layer
from tensorflow.keras import layers
from tensorflow.keras.applications import ResNet50V2,EfficientNetV2S, MobileNetV2, ConvNeXtTiny, ConvNeXtBase


# Strategy for TPU training
regularization_parameter=0.01
dropout_value=0.2
input_shape=(18,96,64,1)

# Create a custom wrapper layer
class CustomKerasLayer(Layer):
    def __init__(self, hub_layer, **kwargs):
        super(CustomKerasLayer, self).__init__(**kwargs)
        self.hub_layer = hub_layer

    def call(self, inputs):
        return self.hub_layer(inputs)
    def compute_output_shape(self, input_shape):
        # Compute the output shape based on the VGGish output
        # Here you need to replace the following with the actual output shape of VGGish
        # Assuming VGGish outputs a fixed shape (for example, (None, 128))
        return (input_shape[0], 128)  # Adjust the second dimension based on VGGish's output
    
with strategy.scope():

    # Load ResNet50V2 with pre-trained ImageNet weights and modify for your task
#     base_model = ResNet50V2(weights='imagenet', include_top=False, input_shape=(128,128, 3))
#     base_model.trainable=True
#     for layer in base_model.layers[:-10]:
#         layer.trainable=False
    base_model = hub.KerasLayer('https://www.kaggle.com/models/google/vggish/TensorFlow2/vggish/1',trainable=True)
#     base_model = EfficientNetV2M(weights='imagenet', include_top=False, input_shape=(256, 256, 3))
    base_model = CustomKerasLayer(base_model)
    # Add custom layers
    inputs = Input(shape=input_shape)
    
    # Apply TimeDistributed to wrap the base CNN
    x = TimeDistributed(base_model)(inputs)

    # Add pooling layer to reduce dimensionality for each time step
    x = TimeDistributed(GlobalAveragePooling2D())(x)
    # Now x has shape (batch_size, time_steps, features)
    
#     x = Flatten()(x)
    # Add LSTM Layer
#     x = LSTM(256, return_sequences=True)(x)  # First LSTM with 256 units
    x = LSTM(128, return_sequences=False)(x)  # LSTM with 64 units
    
    # Increase regularization with more dropout and L2
#     x = Dense(1024, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
# #     x = BatchNormalization()(x)
#     x = Dropout(dropout_value)(x)
    x = Dense(512, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
    x = Dropout(dropout_value)(x)
    x = Dense(256, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
    x = Dropout(dropout_value)(x)
    x = Dense(128, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
    x = Dropout(dropout_value)(x)

    # Output layers for age and gender classification"
    age_output = Dense(3, activation='softmax', name='age')(x)
    gender_output = Dense(1, activation='sigmoid', name='gender')(x)

    # Learning Rate Decay
    reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss',    # Metric to monitor
        factor=0.1,            # Factor by which to reduce the learning rate
        patience=3,            # Number of epochs with no improvement after which the learning rate is reduced
        min_lr=1e-7            # Lower bound on the learning rate
    )
    # Early Stopping 
    early_stopping = EarlyStopping(monitor='val_loss', patience=20, restore_best_weights=True)
    
    # Create the final model
    model = Model(inputs=inputs, outputs=[age_output, gender_output])

    # Compile the model with appropriate loss functions
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
                  loss={'age': 'categorical_crossentropy', 'gender': 'binary_crossentropy'},
                  metrics={'age': 'accuracy', 'gender': 'accuracy'}) ERROR:absl:hub.KerasLayer is trainable but has zero trainable weights.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[18], line 3
      1 # Train the model with data augmentation
      2 with strategy.scope():
----> 3     history = model.fit(train_dataset,
      4                         validation_data=dev_dataset,
      5                         epochs=100, callbacks=[early_stopping,reduce_lr])

File /usr/local/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:122, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    119     filtered_tb = _process_traceback_frames(e.__traceback__)
    120     # To get the full stack trace, call:
    121     # `keras.config.disable_traceback_filtering()`
--> 122     raise e.with_traceback(filtered_tb) from None
    123 finally:
    124     del filtered_tb

File /usr/local/lib/python3.10/site-packages/keras/src/trainers/trainer.py:1040, in Trainer._symbolic_build(self, iterator, data_batch)
   1038     y_pred = backend.compute_output_spec(self, x, training=False)
   1039 except Exception as e:
-> 1040     raise RuntimeError(
   1041         "Unable to automatically build the model. "
   1042         "Please build it yourself before calling "
   1043         "fit/evaluate/predict. "
   1044         "A model is 'built' when its variables have "
   1045         "been created and its `self.built` attribute "
   1046         "is True. Usually, calling the model on a batch "
   1047         "of data is the right way to build it.\n"
   1048         "Exception encountered:\n"
   1049         f"'{e}'"
   1050     )
   1051 if compile_metrics_unbuilt:
   1052     # Build all metric state with `backend.compute_output_spec`.
   1053     backend.compute_output_spec(
   1054         self.compute_metrics,
   1055         x,
   (...)
   1058         sample_weight=sample_weight,
   1059     )

RuntimeError: Unable to automatically build the model. Please build it yourself before calling fit/evaluate/predict. A model is 'built' when its variables have been created and its `self.built` attribute is True. Usually, calling the model on a batch of data is the right way to build it.
Exception encountered:
'Exception encountered when calling layer 'keras_layer_2' (type KerasLayer).

in user code:

    File "/usr/local/lib/python3.10/site-packages/tensorflow_hub/keras_layer.py", line 242, in call  *
        result = f()

    TypeError: Binding inputs to tf.function failed due to `Can not cast TensorSpec(shape=(4, 96, 64, 1), dtype=tf.float32, name=None) to TensorSpec(shape=(None,), dtype=tf.float32, name=None)`. Received args: (<tf.Tensor 'functional_1_1/time_distributed_3_1/strided_slice:0' shape=(4, 96, 64, 1) dtype=float32>,) and kwargs: {} for signature: (waveform: TensorSpec(shape=(None,), dtype=tf.float32, name=None)).


Call arguments received by layer 'keras_layer_2' (type KerasLayer):
  • inputs=tf.Tensor(shape=(4, 96, 64, 1), dtype=float32)
  • training=None'