2
\$\begingroup\$

I use the following code in order to assess the quality of an audio, which is based on this original-project: MOSNet. I call compute_mosnet_score() in a loop with a different audio file in each iteration.

import os
import scipy
import librosa
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import Model, layers
from tensorflow.keras.constraints import max_norm
from tensorflow.keras.layers import Dense, Dropout, Conv2D
from tensorflow.keras.layers import LSTM, TimeDistributed, Bidirectional


class MOSNet():
    def __init__(self, window, wave_handler, hop=None):
        # init
        self.wave_handler = wave_handler

        # constants
        self.fixed_rate = 16000
        self.mono       = True
        self.absolute   = True
        self.FFT_SIZE   = 512
        self.SGRAM_DIM  = self.FFT_SIZE // 2 + 1
        self.HOP_LENGTH = 256
        self.WIN_LENGTH = 512

        _input = keras.Input(shape=(None, 257))
        re_input = layers.Reshape((-1, 257, 1), input_shape=(-1, 257))(_input)

        # CNN
        conv1 = (Conv2D(16, (3, 3), strides=(1, 1), activation='relu', padding='same'))(re_input)
        conv1 = (Conv2D(16, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv1)
        conv1 = (Conv2D(16, (3, 3), strides=(1, 3), activation='relu', padding='same'))(conv1)

        conv2 = (Conv2D(32, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv1)
        conv2 = (Conv2D(32, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv2)
        conv2 = (Conv2D(32, (3, 3), strides=(1, 3), activation='relu', padding='same'))(conv2)

        conv3 = (Conv2D(64, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv2)
        conv3 = (Conv2D(64, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv3)
        conv3 = (Conv2D(64, (3, 3), strides=(1, 3), activation='relu', padding='same'))(conv3)

        conv4 = (Conv2D(128, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv3)
        conv4 = (Conv2D(128, (3, 3), strides=(1, 1), activation='relu', padding='same'))(conv4)
        conv4 = (Conv2D(128, (3, 3), strides=(1, 3), activation='relu', padding='same'))(conv4)

        re_shape = layers.Reshape((-1, 4 * 128), input_shape=(-1, 4, 128))(conv4)

        # BLSTM
        blstm1 = Bidirectional(LSTM(128, return_sequences=True, dropout=0.3,
                                    recurrent_dropout=0.3,
                                    recurrent_constraint=max_norm(0.00001)),
                               merge_mode='concat')(re_shape)

        # DNN
        flatten = TimeDistributed(layers.Flatten())(blstm1)
        dense1 = TimeDistributed(Dense(128, activation='relu'))(flatten)
        dense1 = Dropout(0.3)(dense1)

        frame_score   = TimeDistributed(Dense(1), name='frame')(dense1)
        average_score = layers.GlobalAveragePooling1D(name='avg')(frame_score)

        self.model = Model(outputs=[average_score, frame_score], inputs=_input)

        # weights are in the directory of this file
        pre_trained_dir = "resources"

        # load pre-trained weights. CNN_BLSTM is reported as best
        self.model.load_weights(os.path.join(pre_trained_dir, 'cnn_blstm.h5'))

    def compute_mosnet_score(self, wav_file, fs):
        """
        Evaluate the audio quality using a MOSnet (neural network).
        """
        # compute mosnet score
        sig4mosnet, sig4mosnetlen = self.wave_handler.ffmpeg_load_audio(wav_file, fs)
        mosnet_scores, avg_mos_score = self.predict_mos(sig4mosnet, self.fixed_rate)
        return avg_mos_score

    def predict_wrapper(self, mag):
        res = self.model.predict(mag, verbose=0, batch_size=1, steps=1)
        return res

    def get_mag(self, audios, rate):
        """
        Get signal magnitude.
        """
        # stft. D: (1+n_fft//2, T) -> magnitude spectrogram
        mag = np.abs(librosa.stft(y=np.asfortranarray(audios),
                                  n_fft=self.FFT_SIZE,
                                  hop_length=self.HOP_LENGTH,
                                  win_length=self.WIN_LENGTH,
                                  window=scipy.signal.hamming))  # (1+n_fft/2, T)

        # shape in (T, 1+n_fft/2) -> now call the actual MOSnet
        mag = np.transpose(mag.astype(np.float32))
        return mag[None, ...]

    def predict_mos(self, audios, rate):
        """
        Predict the Mean Objective Score of a given frame.
        """
        mag = self.get_mag(audios, rate)
        res = self.predict_wrapper(mag)
        return res[1][0], round(res[0][0][0], 3)                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      

Since each iteration, the code computes the score using a different audio, it means the input size changes. This causes the following warning:

WARNING:tensorflow:6 out of the last 12 calls to <function _make_execution_function.<locals>.distributed_function at 0x7f999459e170> triggered tf.function retracing. 
Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing.
Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.

Unfortunately, this retracing is slowing down the code so is there any alternatives here that will disable retracing or accelerate the code? Obviously, I can use GPU accelerations, but those are not an option in my case. I already tried to frame the audio signal and feed it to the prediction fucntion, which reduces the warnings but doesn't remove them all. However that is not practical since it add extra calls and slows the code down. I also tried @tf.function and @tf.function( experimental_relax_shapes=True) for def predict_wrapper(self, mag) but it resulted in a ValueError: in converted code.

*I am using :

tensorflow-base           2.1.0           
tensorflow-estimator      2.1.0
\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.