如何在Keras（TensorFlow后端）中用不同尺寸图像训练UNet？

First, let's break down why you're hitting this error: when you append variable-sized images to a list and convert it to a NumPy array, NumPy can't create a proper 4D tensor (since each image has a unique height/width). Instead, it wraps each image as a single element, resulting in a (670, 1) shape—this doesn't match what your input layer Input((None, None, 3)) expects (which is a 4D array of shape (num_samples, height, width, 3)).

Here are three practical solutions tailored to different use cases:

1. Resize All Images to a Fixed Size (Simplest Approach) #

If your task can tolerate losing some aspect ratio information, this is the fastest fix. Just resize every image to a consistent size when loading them:

from skimage.transform import resize  # Or use cv2.resize if you prefer OpenCV
import numpy as np

X_train = []
image_ids = [...]  # Your list of image IDs

for id_ in image_ids:
    img = imread(f"{path}/images/{id_}.png")[:, :, :3]
    # Resize to a fixed size, e.g., 256x256
    resized_img = resize(img, (256, 256), preserve_range=True).astype(np.uint8)
    X_train.append(resized_img)

# Now X_train will be a 4D array of shape (670, 256, 256, 3)
X_train = np.array(X_train)
Y_train = np.array(Y_train)

# Train as normal
model.fit(X_train, Y_train, validation_split=0.1, batch_size=8, epochs=30)

This ensures all inputs have the same shape, eliminating the dimension mismatch error entirely.

2. Use a Custom Generator (Preserve Original Sizes) #

If your model is a fully convolutional network (FCN) (no dense layers with fixed input sizes) and you need to keep the original image dimensions (e.g., for image segmentation or precise object detection), use Keras' Sequence class to build a custom generator:

from keras.utils import Sequence
import numpy as np

class VariableSizeImageGenerator(Sequence):
    def __init__(self, image_ids, labels, img_path, batch_size=8):
        self.image_ids = image_ids
        self.labels = labels
        self.img_path = img_path
        self.batch_size = batch_size

    def __len__(self):
        # Calculate number of batches per epoch
        return int(np.ceil(len(self.image_ids) / self.batch_size))

    def __getitem__(self, idx):
        # Fetch the current batch of IDs and labels
        batch_ids = self.image_ids[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_labels = self.labels[idx*self.batch_size : (idx+1)*self.batch_size]

        # Load images with their original sizes
        batch_imgs = []
        for id_ in batch_ids:
            img = imread(f"{self.img_path}/images/{id_}.png")[:, :, :3]
            batch_imgs.append(img)
        
        # Return batch arrays (Keras handles variable sizes for compatible models)
        return np.array(batch_imgs), np.array(batch_labels)

# Split your data into train/validation sets
train_ids = [...]  # 90% of your image IDs
val_ids = [...]    # 10% for validation
train_labels = [...]
val_labels = [...]

# Create generators
train_gen = VariableSizeImageGenerator(train_ids, train_labels, path, batch_size=8)
val_gen = VariableSizeImageGenerator(val_ids, val_labels, path, batch_size=8)

# Train using the generators
model.fit(train_gen, validation_data=val_gen, epochs=30)

This method preserves your images' original dimensions, but make sure your model doesn't include any dense layers that require fixed input sizes.

3. Use tf.data.Dataset (TensorFlow-Native Flexibility) #

If you're using TensorFlow 2.x, tf.data.Dataset is a powerful, efficient way to handle variable-sized data. It can automatically pad images to the largest size in each batch or keep variable dimensions as-is:

import tensorflow as tf

def load_image(image_id, label):
    # Load and preprocess the image
    img_path = tf.strings.join([path, "/images/", image_id, ".png"])
    img = tf.io.read_file(img_path)
    img = tf.image.decode_png(img, channels=3)
    # Normalize pixel values (optional but recommended)
    img = tf.cast(img, tf.float32) / 255.0
    return img, label

# Create dataset from your IDs and labels
dataset = tf.data.Dataset.from_tensor_slices((image_ids, Y_train))
dataset = dataset.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)

# Split into train/validation
val_size = int(0.1 * len(image_ids))
train_dataset = dataset.skip(val_size).shuffle(100).batch(8)
val_dataset = dataset.take(val_size).batch(8)

# Train the model
model.fit(train_dataset, validation_data=val_dataset, epochs=30)

When using batch(8), TensorFlow automatically pads images in each batch to match the largest height/width in that batch. Your model (with Input((None, None, 3))) will handle this dynamic sizing seamlessly.

Quick Recap #

• Use fixed resizing if you don't need to preserve original image dimensions.
• Use custom generators or tf.data.Dataset if you need to keep variable sizes and have a model compatible with dynamic inputs.

内容的提问来源于stack exchange，提问作者Дмитрий Митюра