TensorFlow Dataset.map传入numpy函数时出现异常传参问题

开发者阿泽

2026-9-2

问题现象

重启Jupyter Notebook后，未做修改的tf.data数据加载代码突然出现异常，此前逻辑可正常运行。
初始通过tf.data.Dataset.from_tensor_slices((specgram_files, labels))构建数据集时，单样本读取结果符合预期：返回元组第一个元素为bytes格式的npy文件路径张量，第二个为对应标签张量。
经tf.numpy_function封装numpy读文件逻辑、传入Dataset.map()做映射后，预期每次函数调用仅接收单个文件路径，实际read_npy_file接收到的是多个无分隔直接拼接的路径字节串，导致np.load()无法正常解析路径报错。
异常传入的参数打印结果如下：

b'./challengeA_data/log_spectrogram/2603ebb3-3cd3-43cc-98ef-0c128c515863.npy'b'./challengeA_data/log_spectrogram/fab6a266-e97a-4935-a0c3-444fc4426fc5.npy'b'./challengeA_data/log_spectrogram/93014682-60a2-45bd-9c9e-7f3c97b83be9.npy'b'./challengeA_data/log_spectrogram/710f2430-5da3-4822-a252-6ad3601b92d9.npy'b'./challengeA_data/log_spectrogram/e757058c-91de-4381-8184-65f001c95647.npy'


b'./challengeA_data/log_spectrogram/38b12689-04ba-422b-a972-5856b05ca868.npy'
b'./challengeA_data/log_spectrogram/7c9ccc04-a2d2-4eec-bafd-0c97b3658c26.npy'b'./challengeA_data/log_spectrogram/c7cc3520-7218-4d07-9f0a-6bd7bb90a551.npy'



b'./challengeA_data/log_spectrogram/21f6060a-9766-4810-bd7c-0437f47ccb98.npy'

可复现异常的核心代码逻辑：

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

specgram_ds = tf.data.Dataset.from_tensor_slices((specgram_files, labels))

specgram_ds = specgram_ds.map(
                    lambda file, label: tuple([tf.numpy_function(read_npy_file, [file], [tf.float32]), label]),
                    num_parallel_calls=tf.data.AUTOTUNE)

异常根因

核心问题是tf.numpy_function封装的自定义函数返回值没有声明明确的静态张量形状。当开启num_parallel_calls=tf.data.AUTOTUNE自动并行优化时，tf.data运行时无法判断单个样本的张量边界，会默认将多个连续样本的输入拼接为连续内存块传入自定义函数，最终出现多个路径字节串直接拼接传入的现象。
重启前代码可正常运行，是因为当时AUTOTUNE根据运行环境状态选择的调度策略未触发批量拼接的优化路径；重启后运行时缓存、资源状态重置，自动调优选择了新的并行调度策略，触发该问题。

修复方案

方案1（最小修改，推荐）：给tf.numpy_function返回的张量补充明确的静态形状声明，让tf.data可以正确识别单样本边界。注意将代码中的SPEC_SHAPE替换为你自己存储的npy频谱数组的实际固定形状，示例代码：

# 替换为实际的npy数组形状，例如(128, 256)对应128个频点、256个时间步的频谱图
SPEC_SHAPE = (128, 128)

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

specgram_ds = specgram_ds.map(
    lambda file, label: (
        tf.numpy_function(read_npy_file, [file], tf.float32).set_shape(SPEC_SHAPE),
        label
    ),
    num_parallel_calls=tf.data.AUTOTUNE
)

方案2（临时验证用）：关闭并行映射，将num_parallel_calls=tf.data.AUTOTUNE修改为num_parallel_calls=1，此时运行时不会做批量拼接处理，可临时解决问题，但数据加载性能会明显下降，不建议生产环境使用。
方案3（最稳妥，规避适配问题）：使用TensorFlow原生IO接口读取文件，完全脱离numpy函数和tf静态图的适配逻辑，注意原生接口解析npy需要固定数组形状，示例代码：

SPEC_SHAPE = (128, 128)

def read_npy_tf(file_path, label):
    npy_bytes = tf.io.read_file(file_path)
    # 解析无压缩np.save存储的float32格式npy文件
    spec = tf.io.decode_raw(npy_bytes, tf.float32)
    spec = tf.reshape(spec, SPEC_SHAPE)
    return spec, label

specgram_ds = specgram_ds.map(read_npy_tf, num_parallel_calls=tf.data.AUTOTUNE)

注意：如果你的npy文件是用np.savez_compressed存储的压缩格式，不要使用方案3，优先选择方案1。

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