如何在TensorFlow中加载不同尺寸图像到内存并使用Dataset API读取

Absolutely! Loading data directly from memory is a fantastic way to speed up your training pipeline when disk I/O becomes a bottleneck. Let’s break down how to implement this, along with handling variable-sized images, using your existing code as a starting point.

There are two main approaches depending on your dataset size:

Option 1: Preload All Images into Numpy Arrays (Best for Small Datasets) #

If your dataset fits comfortably in RAM, you can preload all images into numpy arrays first, then create a tf.data.Dataset directly from these in-memory arrays. This eliminates disk reads entirely after the initial load.

Here’s how to adjust your workflow:

import numpy as np
from PIL import Image
import tensorflow as tf

# Helper to load images into numpy arrays
def load_image_to_numpy(filename):
    with Image.open(filename) as img:
        return np.array(img)

# Your filenames and labels (use regular Python lists here, not tf.constant)
filenames = ["/var/data/image1.jpg", "/var/data/image2.jpg", ...]
labels = [0, 37, ...]

# Preload all images into memory
images_np = np.array([load_image_to_numpy(f) for f in filenames])
labels_np = np.array(labels)

# Create dataset directly from in-memory arrays
dataset = tf.data.Dataset.from_tensor_slices((images_np, labels_np))
# Add your preprocessing logic
dataset = dataset.map(_parse_memory_data, num_parallel_calls=tf.data.AUTOTUNE)

Option 2: Use tf.data.Dataset.cache() (Great for Larger Datasets) #

If your dataset is too big to fit entirely in RAM, use cache() to store processed data in memory after the first epoch. Subsequent epochs will skip disk reads and pull data directly from the cache.

Modify your existing code like this:

import tensorflow as tf

def _parse_function(filename, label):
    # Use tf.io.read_file instead of the deprecated tf.read_file
    image_string = tf.io.read_file(filename)
    # Add expand_animations=False to avoid issues with GIFs
    image_decoded = tf.image.decode_image(image_string, expand_animations=False)
    # Add your preprocessing here (we'll cover variable-sized images next)
    image_processed = tf.image.resize_with_pad(image_decoded, target_height=28, target_width=28)
    image_processed = tf.cast(image_processed, tf.float32) / 255.0  # Normalize
    return image_processed, label

filenames = tf.constant(["/var/data/image1.jpg", "/var/data/image2.jpg", ...])
labels = tf.constant([0, 37, ...])

dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
# Use multi-threading for faster preprocessing
dataset = dataset.map(_parse_function, num_parallel_calls=tf.data.AUTOTUNE)
# Cache processed data in memory (add a path like "./disk_cache" to use disk instead)
dataset = dataset.cache()
# Add shuffle, batch, and prefetch for optimal performance
dataset = dataset.shuffle(buffer_size=1000).batch(32).prefetch(tf.data.AUTOTUNE)

For images with different dimensions, here are the most common preprocessing strategies:

• Resize with aspect ratio preservation: Use tf.image.resize_with_pad to resize images while maintaining their aspect ratio, padding any extra space with black pixels. This avoids distorting your images.
• Force resize to fixed size: Use tf.image.resize if you don’t mind aspect ratio distortion (common in some classification tasks).
• Random crop augmentation: Resize images to a slightly larger size first, then randomly crop to your target dimension for data augmentation:

  def _parse_augment(filename, label):
      image_string = tf.io.read_file(filename)
      image_decoded = tf.image.decode_image(image_string)
      # Resize to a larger intermediate size
      image_resized = tf.image.resize(image_decoded, [32, 32])
      # Randomly crop to the target size
      image_cropped = tf.image.random_crop(image_resized, [28, 28, tf.shape(image_resized)[-1]])
      image_normalized = image_cropped / 255.0
      return image_normalized, label
  

• Pad to maximum dataset size: If you want to keep original image proportions without resizing, calculate the largest height/width in your dataset and pad all images to that size:

  # First calculate max dimensions in Python
  max_height = max(Image.open(f).height for f in filenames)
  max_width = max(Image.open(f).width for f in filenames)
  
  def _parse_pad(filename, label):
      image_string = tf.io.read_file(filename)
      image_decoded = tf.image.decode_image(image_string)
      # Pad to the maximum dimensions
      image_padded = tf.image.pad_to_bounding_box(
          image_decoded,
          offset_height=0,
          offset_width=0,
          target_height=max_height,
          target_width=max_width
      )
      image_normalized = image_padded / 255.0
      return image_normalized, label
  

• Always use num_parallel_calls=tf.data.AUTOTUNE in map() to leverage multi-threaded preprocessing.
• Add prefetch(tf.data.AUTOTUNE) at the end of your dataset pipeline to overlap data preprocessing with model training.
• For the preloaded numpy array approach, check memory usage first with images_np.nbytes to ensure it fits in your RAM.

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