一种基于LSTM的音频质量检测方案

def frame_mse_tf(y_true, y_pred):
    True_pesq = y_true[:,0,:]
    loss = tf.constant(0, dtype=tf.float32)
    for i in range(y_true.shape[0]):
        loss += (10**(True_pesq[i] - 4.5)) * tf.reduce_mean(tf.math.square(y_true[i] - y_pred[i]))
    return loss / tf.constant(y_true.shape[0], dtype=tf.float32)

def train_loop(features, labels1, labels2):
    loss_object = tf.keras.losses.MeanSquaredError()
    with tf.GradientTape() as tape:
        # 预测结果
        predictions1, predictions2 = model(features)
        # 计算 loss
        loss1 = loss_object(labels1, predictions1)
        loss2 = frame_mse_tf(labels2, predictions2)
        loss = loss1 + loss2
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

def read_npy_file(filename):
    data = np.load(filename.numpy().decode())
    return data.astype(np.float32)

def data_preprocessing(feature):
    feature, label1 = feature[...,0] , feature[0][0]
    label2 = label1[0]*np.ones([feature.shape[0],1])
    return feature, label1, label2

def read_feature(filename):
    [feature,] = tf.py_function(read_npy_file,[filename],[tf.float32,])
    data,label1, label2 = tf.py_function(data_preprocessing,[feature],[tf.float32,tf.float32,tf.float32])
    return data, label1, label2

def generate_data(file_path):
    list_ds = tf.data.Dataset.list_files(file_path + '*.npy')
    feature_ds = list_ds.map(read_feature, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return feature_ds

optimizer = tf.keras.optimizers.RMSprop(learning_rate=0.001)
W = model.layers[0].get_weights()
# 读取数据集
ds = generate_data(train_file_path)
# 随机打乱
ds = ds.shuffle(buffer_size=1000).padded_batch(BATCH_SIZE,padded_shapes=([None, None], [None], [None, None])).prefetch(tf.data.experimental.AUTOTUNE)
for step, (x, y, z) in enumerate(ds):
    loss = train_loop(x, y, z)