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
