How to create a TFJS model?

[Mattk70]

Hi @kahst ,

I have trained a custom classifier from BirdNET, and it seems to perform reasonably well. I'd like to import this model into Chirpity, which depends on a TFJS model. I've saved my classifier as a "raven" model, but as I think we discussed, because of the custom layers in BirdNET, it's not straightforward to get tensorflowjs_converter to run. I was hoping you might share the technique you used to convert the TFJS model in this repo, since I have had no success to date - and I have tried all the conversion option permutations I can think of!

Many thanks!

[kahst]

Ah, yes, that's tricky. You need the original Keras model.

Can you try the following workflow (using TF 2.15)?

• download the .h5 file here
• implement the custom spectrogram layer like this:

import tensorflow as tf

class MelSpecLayerSimple(tf.keras.layers.Layer):

    def __init__(self, sample_rate, spec_shape, frame_step, frame_length, fmin, fmax, data_format, **kwargs):
        super(MelSpecLayerSimple, self).__init__(**kwargs)
        self.sample_rate = sample_rate
        self.spec_shape = spec_shape
        self.data_format = data_format
        self.frame_step = frame_step
        self.frame_length = frame_length
        self.fmin=fmin
        self.fmax=fmax

        self.mel_filterbank = tf.signal.linear_to_mel_weight_matrix(
                                num_mel_bins=self.spec_shape[0],
                                num_spectrogram_bins=self.frame_length // 2 + 1,
                                sample_rate=self.sample_rate,
                                lower_edge_hertz=self.fmin,
                                upper_edge_hertz=self.fmax,
                                dtype=tf.float32)  

    def build(self, input_shape):           
        self.mag_scale = self.add_weight(name='magnitude_scaling', 
                                         initializer=tf.keras.initializers.Constant(value=1.23),
                                         trainable=True)                                          
        super(MelSpecLayerSimple, self).build(input_shape)        

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_last':
            return tf.TensorShape((None, self.spec_shape[0], self.spec_shape[1], 1))
        else:
            return tf.TensorShape((None, 1, self.spec_shape[0], self.spec_shape[1]))

    def call(self, inputs, training=None):

        # Normalize values between -1 and 1
        inputs = tf.math.subtract(inputs, tf.math.reduce_min(inputs, axis=1, keepdims=True))
        inputs = tf.math.divide(inputs, tf.math.reduce_max(inputs, axis=1, keepdims=True) + 0.000001)
        inputs = tf.math.subtract(inputs, 0.5)
        inputs = tf.math.multiply(inputs, 2.0)        

        # Perform STFT    
        spec = tf.signal.stft(inputs,
                              self.frame_length,
                              self.frame_step,
                              fft_length=self.frame_length,
                              window_fn=tf.signal.hann_window,
                              pad_end=False,
                              name='stft')    

        # Cast from complex to float
        spec = tf.dtypes.cast(spec, tf.float32)

        # Apply mel scale              
        spec = tf.tensordot(spec, self.mel_filterbank, 1)
        
        # Convert to power spectrogram
        spec = tf.math.pow(spec, 2.0)

        # Convert magnitudes using nonlinearity
        spec = tf.math.pow(spec, 1.0 / (1.0 + tf.math.exp(self.mag_scale)))

        # Flip spec horizontally
        spec = tf.reverse(spec, axis=[2])

        # Swap axes to fit input shape
        spec = tf.transpose(spec, [0, 2, 1])

        # Add channel axis        
        if self.data_format == 'channels_last':
            spec = tf.expand_dims(spec, -1)
        else:
            spec = tf.expand_dims(spec, 1)      

        return spec

    def get_config(self):
        config = {'data_format': self.data_format,
                  'sample_rate': self.sample_rate,
                  'spec_shape': self.spec_shape,
                  'frame_step': self.frame_step,
                  'fmin': self.fmin,
                  'fmax': self.fmax,
                  'frame_length': self.frame_length,
                  'mel_filterbank': self.mel_filterbank.numpy().tolist()
                  }
        base_config = super(MelSpecLayerSimple, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

• convert the custom Raven model to TFJS (kind of weird, but it seems to work)

import tensorflowjs as tfjs

def raven2tfjs(raven_model_path, keras_model_path, output_model_path):  

    # Load Raven model (== Tensorflow saved model)
    raven_model = tf.saved_model.load(raven_model_path)
    weights = [var for var in raven_model.signatures['basic'].variables if 'dense/kernel' in var.name]
    biases = [var for var in raven_model.signatures['basic'].variables if 'dense/bias' in var.name]

    # Load keras model
    keras_model = tf.keras.models.load_model(keras_model_path, custom_objects={'MelSpecLayerSimple': MelSpecLayerSimple})

    # Make new keras model without the classification layer
    # and add a dense layer with num_units = num_classes in Raven model
    new_keras_model = tf.keras.Sequential([tf.keras.Model(inputs=keras_model.input, outputs=keras_model.layers[-3].output)])
    new_keras_model.add(tf.keras.layers.Dense(weights[0].shape[1]))

    # Set weights of classification layer
    new_keras_model.layers[-1].set_weights([weights[0].numpy(), biases[0].numpy()])

    # Save as tfjs model
    tfjs.converters.save_keras_model(new_keras_model, output_model_path)

raven2tfjs('your_custom_analyzer_model', 'BirdNET_GLOBAL_6K_V2.4_Model_Keras.h5', 'Custom_Model.tfjs')

• load the TFJS model as usual in JavaScript (+ new labels file)
• confirm the scores are what you would expect

Can you give it a try? Let me know if it worked.

Best,
Stefan

[Mattk70]

Thanks so much Stefan - this is great. I'm going to be away for a week or so, and won't get around to trying this out 'til I'm back. But, I will get back to you on how I get on - I really appreciate the detailed response.

Happy Easter!

[Mattk70]

I should mention, this #285 (comment) was a super-helpful contribution. I'd had the same memory issue with my dataset before making the changes suggested - which worked like a charm!

[tphakala]

@kahst is range filter available as Keras model? I am trying to convert this into ONNX model

[Mattk70]

Hi Tomi, did you attempt a tflite conversion? From https://github.com/onnx/tensorflow-onnx you can convert a tflite file directly:

python -m tf2onnx.convert --opset 16 --tflite tflite--file --output model.onnx

[tphakala]

No I haven't tested this yet, I have been thinking of it though. Maybe I'll give it a try.

[Mattk70]

@kahst - an update. I have followed your suggestions here and can report that the conversion completes without error (yay!), and I have hacked the custom model into a branch of Chirpity. I applied a tf.sigmoid() function to the logits this converted model produces. The results are similar to the tf.saved_model (raven) output but not as similar as I think they should be. This time, I don't think it is the sample rate issue I had before with the TFJS models. I'm not sure why yet, but will keep digging.

[Mattk70]

Did some more digging, found the problem. Yes, this works! Thanks!