How to create a time-frequency loss function use tensorflow or pytorch

Question

I am doing a project about speech enhancement. I want to use time-frequency loss function like so:

# create a frequency loss function for the speech enhancement with tensorflow
def freq_loss(y_true, y_pred):
    # convert the tensors to numpy arrays
    y_true = y_true.numpy()
    y_pred = y_pred.numpy()
    # calculate the frequency loss
    loss = np.sum(np.abs(np.fft.fft(y_true) - np.fft.fft(y_pred)))
    # return the loss
    return loss

# create a time loss function for the speech enhancement with tensorflow
def time_loss(y_true, y_pred):
    # convert the tensors to numpy arrays
    y_true = y_true.numpy()
    y_pred = y_pred.numpy()
    # calculate the time loss
    loss = np.sum(np.abs(y_true - y_pred))
    # return the loss
    return loss

# create a combined loss function for the speech enhancement with tensorflow
def combined_loss(y_true, y_pred):
    # convert the tensors to numpy arrays
    y_true = y_true.numpy()
    y_pred = y_pred.numpy()
    # calculate the frequency loss
    freq_loss = np.sum(np.abs(np.fft.fft(y_true) - np.fft.fft(y_pred)))

to evaluate the performance of my algorithm about original noisy voice data and noise-supressed voice data.

I have found some papers about this topic, but I can't find great useful code. Can anyone help me?

Answer 1

Given by the open-endedness of the question, here is one of the loss functions you can use between your speech signals

import torch

def loss(y_pred, y_target):
    # Compute the magnitude
    y_pred_spectrum = torch.abs(torch.rfft(y_pred, 1))
    y_target_spectrum = torch.abs(torch.rfft(y_target, 1))

    # get difference between the spectra
    spectra_difference = y_pred_spectrum - y_target_spectrum

    # Apply a weighting function
    weighting_function = (you can add here your criteria for weighting)
    weighted_difference = spectra_difference * weighting_function

    # Compute the loss
    loss = torch.sum(weighted_difference**2)

    return loss