求助：如何用TensorFlow/PyTorch构建语音增强的时频损失函数

原代码存在两个核心问题：一是依赖numpy操作TensorFlow张量，在图模式训练时会抛出错误；二是频域损失的设计不够贴合语音增强的实际需求（直接全局FFT不如短时傅里叶变换STFT更能捕捉语音的时频特性）。以下是修正后的完整实现：

1. 基础工具函数：STFT转换 #

先定义统一的STFT处理函数，确保输入输出都是TensorFlow张量：

import tensorflow as tf

def stft_transform(signal, frame_length=512, frame_step=128, fft_length=512):
    # 计算STFT幅度谱
    stft = tf.signal.stft(
        signals=signal,
        frame_length=frame_length,
        frame_step=frame_step,
        fft_length=fft_length
    )
    return tf.abs(stft)  # 返回幅度谱，忽略相位

2. 时域损失函数 #

用TensorFlow原生操作替代numpy，避免图模式报错：

def time_domain_loss(y_true, y_pred):
    # L1损失，也可以换成L2损失(tf.square后求和)
    return tf.reduce_sum(tf.abs(y_true - y_pred))

3. 频域损失函数 #

改用STFT幅度谱计算，更符合语音增强的评估逻辑：

def freq_domain_loss(y_true, y_pred):
    # 转换为STFT幅度谱
    true_spec = stft_transform(y_true)
    pred_spec = stft_transform(y_pred)
    # 计算幅度谱的L1损失，也可根据需求换成对数幅度谱损失
    return tf.reduce_sum(tf.abs(true_spec - pred_spec))

4. 组合时频损失函数 #

加入可调节的权重参数，平衡时域和频域的损失贡献：

def combined_time_freq_loss(y_true, y_pred, alpha=0.3, beta=0.7):
    """
    alpha: 时域损失的权重
    beta: 频域损失的权重
    注意alpha + beta = 1.0，可根据实验调整
    """
    time_loss = time_domain_loss(y_true, y_pred)
    freq_loss = freq_domain_loss(y_true, y_pred)
    # 归一化损失（可选，避免某一维度损失值过大主导训练）
    time_loss = tf.math.divide(time_loss, tf.reduce_max(time_loss))
    freq_loss = tf.math.divide(freq_loss, tf.reduce_max(freq_loss))
    return alpha * time_loss + beta * freq_loss

关键修正点说明 #

• 全程使用TensorFlow原生API（tf.signal.stft、tf.reduce_sum等），确保在图模式和eager模式下都能正常运行。
• 用STFT替代全局FFT：语音是短时平稳信号，STFT能更好地捕捉不同时间段的频率特征，比全局FFT更适合语音增强任务。
• 加入损失归一化：避免时域和频域损失的量级差异过大，导致模型偏向优化某一方。
• 可扩展：如果需要更贴合听觉感知的损失，可将幅度谱转换为对数幅度谱（tf.math.log(true_spec + 1e-8)），或者加入梅尔谱损失（tf.signal.mfccs_from_log_mel_spectrograms）。

内容的提问来源于stack exchange，提问作者Regan Qing