基于tf.data的多输出模型训练报错：标签与Logits形状不兼容

Let's break down why you're hitting this shape mismatch error and fix it step by step:

The Root Cause #

Your error logits and labels must have the same shape ((None, 1) vs (None, 2)) happens because:

• Your model has two separate output branches: gender_output (shape (None,1) for sigmoid binary classification) and age_output (shape (None,1) for regression).
• But your tf.data.Dataset is returning labels as a single combined tensor of shape (2,), rather than two independent tensors matching each output branch. TensorFlow can't map this single combined tensor to the two model outputs correctly, leading to the shape mismatch for the gender loss calculation.

Step 1: Adjust Label Return Format in process_path #

Instead of returning a single combined label tensor, return two separate label tensors (gender and age) wrapped in a tuple. This matches the model's multi-output structure:

def process_path(file):
    file = file.numpy()
    # Fix: Get full file path (original code had a bug here - list_files returns filenames, not full paths)
    file_path = os.path.join(training_dir, str(bytes.decode(file)))
    # Extract filename from path
    file_name = os.path.basename(file_path)
    gender, age = get_label(file_name) # Unpack gender and age
    
    # Load and preprocess image
    img = tf.io.read_file(file_path)
    img = decode_img(img)
    img = img / 255.0
    
    # Convert labels to tensors with shapes matching model outputs
    gender = tf.convert_to_tensor([gender], dtype=tf.int32) # Shape (1,) to match (None,1) output
    age = tf.convert_to_tensor([age], dtype=tf.int32) # Shape (1,) to match (None,1) output
    
    # Return (image, (gender_label, age_label))
    return img, (gender, age)

Step 2: Update tf.py_function Output Type Specification #

Since we're now returning two label tensors instead of one, we need to update the type argument in tf.py_function:

train_gen = train_gen.map(
    lambda x: tf.py_function(process_path, [x], [tf.float32, (tf.int32, tf.int32)]),
    num_parallel_calls=AUTOTUNE
)
valid_gen = valid_gen.map(
    lambda x: tf.py_function(process_path, [x], [tf.float32, (tf.int32, tf.int32)]),
    num_parallel_calls=AUTOTUNE
)

Step 3: Fix the _set_shapes Function #

Update this function to handle the two separate label tensors and set their shapes correctly:

def _set_shapes(img, labels):
    gender_label, age_label = labels
    # Set image shape
    img.set_shape((128, 128, 3))
    # Set label shapes to match model output shapes
    gender_label.set_shape((1,))
    age_label.set_shape((1,))
    return img, (gender_label, age_label)

Step 4: Verify Dataset Output Shape #

After making these changes, check your dataset's output shape with:

print(train_gen.element_spec)

You should see something like:

(TensorSpec(shape=(128, 128, 3), dtype=tf.float32, name=None),
 (TensorSpec(shape=(1,), dtype=tf.int32, name=None),
  TensorSpec(shape=(1,), dtype=tf.int32, name=None)))

This confirms the dataset is returning the right structure for your multi-output model.

Bonus: Fix list_files Path Issue #

Your original list_files call uses os.listdir(training_dir), which returns filenames instead of full paths. This will cause tf.io.read_file to fail unless your working directory is the training directory. Fix it by using a glob pattern instead:

train_gen = tf.data.Dataset.list_files(os.path.join(training_dir, "*"), shuffle=False)
valid_gen = tf.data.Dataset.list_files(os.path.join(validation_dir, "*"), shuffle=False)

After applying these changes, your model should correctly map each label tensor to its corresponding output branch, and the shape mismatch error should be resolved. The MirroredStrategy will handle multi-GPU training seamlessly with this corrected dataset structure.

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