如何用TFRecords数据集训练TF.Keras模型？输入不匹配问题解决

Let's break down the issue and fix it step by step:

The Root Cause #

Your error happens because your dataset is returning three separate feature tensors (one for each of weight_pounds, gestation_weeks, plurality), but your sequential Keras model expects a single input tensor with shape (3,) (a 3-dimensional feature vector). Additionally, there's a small mistake in how you're parsing the TFRecords.

Fixes to Apply #

We need two key changes: correct the TFRecord parsing function, and merge your separate features into a single input tensor that matches your model's expectations.

1. Fix the TFRecord Parsing Function

You're using tf.io.parse_example, which is designed to parse batches of examples. But TFRecordDataset yields individual example protobuffers one at a time, so you should use tf.io.parse_single_example instead.

2. Merge Features into a Single Tensor

Modify your normalize function to stack the three normalized features into a single 3-dimensional vector, which will align with your model's input shape after batching.

Corrected Full Code #

First, the dataset processing pipeline:

import tensorflow as tf
from tensorflow import keras

# Load the TFRecord dataset
dataset = tf.data.TFRecordDataset("dataset.tfrecords")

# Fixed parsing function for single examples
def _parse_function(proto):
    keys_to_features = {
        'weight_pounds': tf.io.FixedLenFeature([], tf.float32),
        'gestation_weeks': tf.io.FixedLenFeature([], tf.float32),
        'plurality': tf.io.FixedLenFeature([], tf.float32),
        'isMale': tf.io.FixedLenFeature([], tf.float32),
    }
    # Parse a single example protobuffer
    parsed_features = tf.io.parse_single_example(proto, keys_to_features)
    return parsed_features

# Precomputed mean and std for normalization
hold_meanstd = {
    'weight_pounds': [7.234738, 1.330294],
    'gestation_weeks': [38.346464, 4.153269],
    'plurality': [1.035285, 0.196870]
}

def normalize(example):
    # Apply normalization to each feature
    example['weight_pounds'] = (example['weight_pounds'] - hold_meanstd['weight_pounds'][0]) / hold_meanstd['weight_pounds'][1]
    example['gestation_weeks'] = (example['gestation_weeks'] - hold_meanstd['gestation_weeks'][0]) / hold_meanstd['gestation_weeks'][1]
    example['plurality'] = (example['plurality'] - hold_meanstd['plurality'][0]) / hold_meanstd['plurality'][1]
    
    # Extract the label
    label = example.pop('isMale')
    
    # Stack the three features into a single (3,) tensor
    # After batching, this becomes (batch_size, 3), matching the model's input shape
    features = tf.stack([example['weight_pounds'], example['gestation_weeks'], example['plurality']], axis=-1)
    
    return (features, label)

# Rebuild the dataset pipeline
dataset = tf.data.TFRecordDataset(["dataset.tfrecords"]).map(_parse_function)
dataset = dataset.map(normalize)
dataset = dataset.batch(64)

Then the model training code (with a minor optimizer update for compatibility):

Dense = keras.layers.Dense
model = keras.Sequential(
    [
        Dense(500, activation="relu", kernel_initializer='uniform', input_shape=(3,)),
        Dense(200, activation="relu"),
        Dense(100, activation="relu"),
        Dense(25, activation="relu"),
        Dense(1, activation="sigmoid")
    ])

# Use `learning_rate` instead of `lr` for newer TensorFlow versions (backward compatible)
optimizer = keras.optimizers.RMSprop(learning_rate=0.01)
model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=[tf.keras.metrics.AUC()])

# Now this should run without input shape errors
model.fit(dataset)

Why This Works #

• parse_single_example: Correctly parses individual examples from the TFRecord stream, ensuring your features are loaded as scalar tensors instead of batched tensors.
• tf.stack: Combines your three separate scalar features into a single 3-dimensional vector. When you batch the dataset with batch(64), this becomes a tensor of shape (64, 3), which perfectly matches the input_shape=(3,) defined in your first dense layer.

If you ever need to handle multiple distinct inputs (e.g., different types of features), you could use Keras' functional API to define a model with multiple input layers, but for this case, merging into a single vector is the simplest and most efficient solution.

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