应用Meijering滤波器后TensorFlow数据集保存速度极慢求助

环境与背景 #

使用Python 3.8、TensorFlow 2.11.1，对TensorFlow数据集应用Meijering滤波后，调用save()写入磁盘的时间从原本约1分钟骤增至近2小时。

相关代码 #

import tensorflow as tf
from skimage import filters  

# Filtering function
def meijering_filter(x):
    filtered = filters.meijering(x)
    return filtered

# Import training data
training_dataset = tf.keras.utils.image_dataset_from_directory(
        "path_to_training_dataset",
        labels=None,
        batch_size=5,
        image_size=(480, 640),
        shuffle=True,
        seed=42,
        subset='training',
        validation_split=0.2,
        color_mode='grayscale'
    )

normalization_layer = tf.keras.layers.Rescaling(1./255)
normalized_train_dataset = training_dataset.map(lambda x: (normalization_layer(x)))
feat_training_dataset = normalized_train_dataset.map(lambda x: tf.numpy_function(meijering_filter, [x], tf.float32))

# Reshaping data, since the numpy_function() returns tensors with an unknown shape
data_reshape = tf.keras.Sequential([tf.keras.layers.Input(shape=(480, 640, 1))])
feat_training_dataset = feat_training_dataset.map(lambda x: (data_reshape(x)))

# Saving tensorflow dataset for later consumption
feat_training_dataset.save("save_path_on_disk")

已知测试信息 #

• save()需完成一次完整数据集计算
• Meijering滤波器存在计算量，但take(1)测试单批次耗时0.12秒，按1500张图像规模预计仅需数分钟
• 移除data_reshape()操作后耗时无明显变化

耗时剧增的核心原因 #

1. tf.numpy_function的上下文切换开销：该函数会在TensorFlow计算图与Python/NumPy环境之间频繁切换，每处理一个样本就要完成一次跨上下文转换，当数据集规模较大时，这种切换的累积耗时会远超过滤波器本身的计算时间。
2. 单线程+非批量处理：默认的map操作是单线程执行，且skimage.filters.meijering仅针对单张图像处理，没有利用TensorFlow的并行计算能力，也未对批次数据做向量化优化，进一步拖慢了处理速度。
3. 序列化校验额外开销：经过tf.numpy_function处理后的张量，在save()序列化时会触发额外的形状、类型校验，加上单线程处理，拉长了整体保存时间。

可行优化方案 #

1. 替换为TensorFlow原生/兼容的滤波实现 #

优先寻找Meijering滤波器的TensorFlow原生实现，或用tf.image模块中的类似滤波逻辑、自定义TF算子替代skimage函数。这样能完全在TensorFlow计算图内运行，自动利用GPU加速和并行处理，彻底避免上下文切换开销。

2. 批量处理+开启并行map #

修改滤波函数直接处理整批图像，减少上下文切换次数；同时在map操作中设置num_parallel_calls开启并行：

def meijering_filter_batch(batch_x):
    # 处理形状为 (batch_size, 480, 640, 1) 的批次数据
    batch_filtered = []
    for img in batch_x:
        # 去除通道维度以适配skimage的2D输入要求
        img_2d = tf.squeeze(img, axis=-1).numpy()
        filtered = filters.meijering(img_2d)
        # 重新添加通道维度
        batch_filtered.append(tf.expand_dims(filtered, axis=-1))
    return tf.stack(batch_filtered)

# 修改map操作，启用并行处理
feat_training_dataset = normalized_train_dataset.map(
    lambda x: tf.numpy_function(meijering_filter_batch, [x], tf.float32),
    num_parallel_calls=tf.data.AUTOTUNE
)

3. 开启tf.data全流程优化 #

在所有map操作中添加num_parallel_calls=tf.data.AUTOTUNE，让TensorFlow自动适配系统资源调整并行线程数；同时添加prefetch实现预加载：

# 归一化阶段开启并行
normalized_train_dataset = training_dataset.map(
    lambda x: normalization_layer(x),
    num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)

# 滤波处理后也添加预加载
feat_training_dataset = feat_training_dataset.prefetch(tf.data.AUTOTUNE)

4. 先预处理为numpy文件再转TF数据集 #

如果上述优化效果有限，可以先在Python环境中完成所有图像的预处理，保存为numpy文件后再构建TensorFlow数据集：

import numpy as np

all_processed_data = []
for batch in normalized_train_dataset:
    filtered_batch = meijering_filter_batch(batch)
    all_processed_data.append(filtered_batch.numpy())

# 合并所有批次数据
all_processed_data = np.concatenate(all_processed_data, axis=0)
# 保存为numpy文件
np.save("preprocessed_training_data.npy", all_processed_data)

# 后续加载时直接转为TF数据集
loaded_dataset = tf.data.Dataset.from_tensor_slices(all_processed_data)
loaded_dataset.save("save_path_on_disk")

这种方式避免了TensorFlow上下文切换的开销，预处理和保存速度都会大幅提升。

5. 确认GPU加速是否生效 #

用tf.config.list_physical_devices('GPU')检查TensorFlow是否识别并启用了GPU。skimage的操作默认在CPU执行，若能将滤波逻辑迁移到GPU（如用TensorFlow实现），计算速度会有质的提升。

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