Problem incorporating SpecAugument in the training process

Hi,

I'm trying to add a SpecAug layer in the training process of a CNN using the code below:


CLIP_DURATION = 5 
SAMPLING_RATE = 41000
NUM_CHANNELS = 1

INPUT_SHAPE = ((CLIP_DURATION * SAMPLING_RATE), NUM_CHANNELS)

melgram = get_melspectrogram_layer(input_shape = INPUT_SHAPE,
                          n_fft = 2048,
                          hop_length = 512,
                          return_decibel=True,
                          n_mels = 40,
                          mel_f_min = 500,
                          mel_f_max = 15000,
                          input_data_format='channels_last',
                          output_data_format='channels_last')

spec_augment = SpecAugment(freq_mask_param=5,
                          time_mask_param=10,
                          n_freq_masks=2,
                          n_time_masks=3,
                          mask_value=-100)   

model = Sequential()
model.add(melgram)
model.add(spec_augment)

The CNN summary looks like this:

Model: "sequential_2"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 melspectrogram (Sequential)  (None, 397, 40, 1)       0         

 spec_augment_1 (SpecAugment  (None, 397, 40, 1)       0         
 )                                                               

=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________

Compiling and fitting the model

model.compile(loss = 'sparse_categorical_crossentropy', optimizer='adam', metrics = 'accuracy')

early_stop = EarlyStopping(monitor='loss', patience=5)

reduce_LR = ReduceLROnPlateau(monitor="val_loss",factor=0.1,patience=4)

checkpointer = ModelCheckpoint(filepath = 'saved_models/bird_song_classification.hdf5')

model.fit(X_train, y_train, validation_data = (X_val, y_val), epochs = 50, batch_size = 32, callbacks = [early_stop, checkpointer, reduce_LR])

Then I get the following error:

Epoch 1/50
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-35-e58a056ab523>](https://localhost:8080/#) in <module>
      7 checkpointer = ModelCheckpoint(filepath = 'saved_models/bird_song_classification.hdf5')
      8 
----> 9 model.fit(X_train, y_train, validation_data = (X_val, y_val), epochs = 50, batch_size = 32, callbacks = [early_stop, checkpointer, reduce_LR])

6 frames
[/usr/local/lib/python3.7/dist-packages/kapre/augmentation.py](https://localhost:8080/#) in tf___apply_masks_to_axis(self, x, axis, mask_param, n_masks)
     78                 try:
     79                     do_return = True
---> 80                     retval_ = ag__.converted_call(ag__.ld(tf).where, (ag__.ld(mask), ag__.ld(self).mask_value, ag__.ld(x)), None, fscope)
     81                 except:
     82                     do_return = False

TypeError: in user code:

    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1051, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1040, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1030, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 889, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_filepzvfxhgz.py", line 63, in tf__call
        ag__.if_stmt((ag__.ld(training) in (None, False)), if_body_2, else_body_2, get_state_2, set_state_2, ('do_return', 'retval_'), 2)
    File "/tmp/__autograph_generated_filepzvfxhgz.py", line 58, in else_body_2
        retval_ = ag__.converted_call(ag__.ld(tf).map_fn, (), dict(elems=ag__.ld(x), fn=ag__.ld(self)._apply_spec_augment, dtype=ag__.ld(tf).float32, fn_output_signature=ag__.ld(tf).float32), fscope)
    File "/tmp/__autograph_generated_filef27o6c1f.py", line 44, in tf___apply_spec_augment
        ag__.if_stmt((ag__.ld(self).n_time_masks >= 1), if_body_1, else_body_1, get_state_1, set_state_1, ('x',), 1)
    File "/tmp/__autograph_generated_filef27o6c1f.py", line 39, in if_body_1
        x = ag__.converted_call(ag__.ld(self)._apply_masks_to_axis, (ag__.ld(x),), dict(axis=ag__.ld(time_axis), mask_param=ag__.ld(self).time_mask_param, n_masks=ag__.ld(self).n_time_masks), fscope)
    File "/tmp/__autograph_generated_file3vip8w4x.py", line 80, in tf___apply_masks_to_axis
        retval_ = ag__.converted_call(ag__.ld(tf).where, (ag__.ld(mask), ag__.ld(self).mask_value, ag__.ld(x)), None, fscope)

    TypeError: Exception encountered when calling layer "spec_augment_1" (type SpecAugment).

    in user code:

        File "/usr/local/lib/python3.7/dist-packages/kapre/augmentation.py", line 299, in call  *
            elems=x, fn=self._apply_spec_augment, dtype=tf.float32, fn_output_signature=tf.float32
        File "/usr/local/lib/python3.7/dist-packages/kapre/augmentation.py", line 273, in _apply_spec_augment  *
            x = self._apply_masks_to_axis(
        File "/usr/local/lib/python3.7/dist-packages/kapre/augmentation.py", line 254, in _apply_masks_to_axis  *
            return tf.where(mask, self.mask_value, x)

        TypeError: Input 'e' of 'SelectV2' Op has type float32 that does not match type int32 of argument 't'.

    Call arguments received by layer "spec_augment_1" (type SpecAugment):
      • x=tf.Tensor(shape=(None, 397, 40, 1), dtype=float32)
      • training=True
      • kwargs=<class 'inspect._empty'>

The shape of X_train is

(2182, 205000, 1)

I'm using Tensorflow 2.9.2, and Python 3.7.15

When I remove the SpecAug layer everything runs fine. I've tested using only the melspec + a mobile net at the end and it runs smooth. The problem is apparently related to SpecAug layer.

Do you have any idea what could be going wrong here? I appreciate any guidance related to the problem. Best regards.

keunwoochoi / kapre

Problem incorporating SpecAugument in the training process #140