如何实现TFHub YAMNet预训练模型的批量预测并提取嵌入向量

问题原因 #

YAMNet预训练模型的默认调用签名仅支持单个1D音频张量输入（形状为(采样点数量,)），直接传入2D批量数组（形状为(样本数, 采样点数量)）会触发维度不匹配报错。

可行解法 #

方法1：使用tf.map_fn批量遍历推理

适合样本长度不一致的场景，灵活度高：

import tensorflow as tf
import tensorflow_hub as hub
import numpy as np

# 原输入定义
X = np.random.uniform(-1,1,(10, 1000))
yamnet_model_handle = 'https://tfhub.dev/google/yamnet/1'
yamnet_model = hub.load(yamnet_model_handle)

# 单样本推理函数
def single_infer(audio):
    return yamnet_model(audio)

# 批量执行
batch_results = tf.map_fn(
    single_infer,
    tf.constant(X, dtype=tf.float32), # 输入转float32张量避免类型错误
    fn_output_signature=(
        tf.TensorSpec(shape=(None, 521), dtype=tf.float32), # 分类得分输出
        tf.TensorSpec(shape=(None, 1024), dtype=tf.float32), # 嵌入向量输出
        tf.TensorSpec(shape=(None, 64), dtype=tf.float32), # 语谱图输出
    )
)

all_scores, all_embeddings, all_spectrograms = batch_results

方法2：封装支持批量输入的Keras层

适合所有样本长度一致的场景，推理效率更高：

yamnet_batch_layer = hub.KerasLayer(
    yamnet_model_handle,
    trainable=False,
    # 显式指定批量输入签名，None对应动态batch大小
    input_signature=tf.TensorSpec(shape=(None, 1000), dtype=tf.float32)
)

# 直接传入批量输入即可得到结果
all_scores, all_embeddings, all_spectrograms = yamnet_batch_layer(tf.constant(X, dtype=tf.float32))

# 如果仅需要嵌入向量，可在定义层时添加output_key参数直接提取：
# yamnet_emb_layer = hub.KerasLayer(
#     yamnet_model_handle,
#     trainable=False,
#     output_key='embeddings',
#     input_signature=tf.TensorSpec(shape=(None, 1000), dtype=tf.float32)
# )
# all_embeddings = yamnet_emb_layer(tf.constant(X, dtype=tf.float32))

注意事项 #

• 输入需提前转为tf.float32类型的张量，避免numpy数组类型不匹配导致的报错
• 若单条音频推理返回多帧嵌入向量，可按需求取均值、最大值或直接拼接得到单条音频的全局嵌入

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