将HDF5转tf.data.Dataset训练自编码器时遇维度错误求助

错误原因 #

1. tf.py_function返回张量形状未明确：tf.py_function不会自动保留输入输出张量的形状信息，你从HDF5读取的epsilon（形状(128,128,1)）和field（形状(128,128,6)）经过tf.py_function后，张量形状变为未知状态。当执行batch操作后，TensorFlow无法正确推断最终输入模型的张量维度，而自编码器期望接收4维输入（批量大小+128+128+1），因此抛出维度不匹配的错误。
2. 函数名笔误：你的代码中map调用的是load_files，但实际定义的加载函数是load_file，虽然遍历数据集时输出正常，但这可能导致逻辑混乱，需要统一函数名。

解决步骤 #

1. 修正函数名并明确张量形状 #

修改wrapper函数，为返回的张量显式设置形状，同时修正map中的函数调用：

def load_file(file):
    hf = h5py.File(file.numpy(),'r')
    epsilon = np.array(hf.get('epsilon')) # 形状(128,128,1)
    field = np.array(hf.get('field')) # 形状(128,128,6)
    hf.close()
    return epsilon, field

def wrapper(file):
    e, f = tf.py_function(load_file, [file], (tf.float64, tf.float64))
    # 显式设置每个样本的形状
    e.set_shape((128, 128, 1))
    f.set_shape((128, 128, 6))
    return e, f

# 修改map调用，使用wrapper函数
dataset = (dataset
    .shuffle(1024)
    .map(wrapper, num_parallel_calls=AUTOTUNE)
    .cache()
    .repeat()
    .batch(BS)
    .prefetch(AUTOTUNE)
)

2. 确认模型输入层定义 #

确保自编码器的输入层明确指定输入形状为(128, 128, 1)，示例如下：

input_layer = tf.keras.layers.Input(shape=(128, 128, 1))
# 后续自编码器层定义...

3. 可选：验证数据集输出形状 #

在训练前可以添加代码验证数据集的输出形状，确保符合预期：

for epsilon_batch, field_batch in dataset.take(1):
    print("Epsilon batch shape:", epsilon_batch.shape) # 应输出(BS, 128, 128, 1)
    print("Field batch shape:", field_batch.shape) # 应输出(BS, 128, 128, 6)

内容的提问来源于stack exchange，提问作者münsteraner