基于TensorFlow GPU-Jupyter镜像构建Dockerfile还是重构图像加载逻辑？

问题背景 #

在CPU环境使用TensorFlow 2.6.2时，可通过tf.keras.preprocessing.image_dataset_from_directory轻松加载按类别划分的图像数据集，但切换到官方tensorflow/tensorflow:latest-gpu-py3-jupyter镜像后，发现镜像内置TF版本为2.1.0，该API尚未推出；尝试用load_img读取目录时又报错IsADirectoryError，因为load_img仅支持单张图片文件。

你的数据集目录结构如下：

my_folder
│
└───Training
│   └───Class_A
│       01234.jpg
│       56789.jpg
│       ...
│   └───Class_B
│       01234.jpg
│       56789.jpg
│       ...
└───Validation
│   └───Class_A
│       01234.jpg
│       56789.jpg
│       ...
│   └───Class_B
│       01234.jpg
│       56789.jpg
│       ...

方案对比与实现 #

方案一：升级TensorFlow GPU-Jupyter镜像（推荐） #

image_dataset_from_directory是TensorFlow 2.3及以上版本才有的API，直接使用新版本官方镜像即可复用原有代码，同时享受新版本的GPU优化。

实现方式：

1. 使用最新稳定版TF GPU-Jupyter镜像（或指定具体版本，如2.15.0-gpu-py3-jupyter）：

   docker run -it --gpus all -p 8888:8888 -v /本地路径/my_folder:/tf/my_folder tensorflow/tensorflow:latest-gpu-py3-jupyter
   

   其中/本地路径/my_folder替换为你本地数据集的绝对路径，挂载到容器内的/tf/my_folder目录。

2. 容器启动后，直接复用原有代码加载数据：

   import tensorflow as tf
   
   path_training = 'my_folder/Training/'
   path_validation = 'my_folder/Validation/'
   image_size = (90, 90)
   
   train_ds = tf.keras.preprocessing.image_dataset_from_directory(
       path_training,
       seed=1993, 
       image_size=image_size
   )
   
   val_ds = tf.keras.preprocessing.image_dataset_from_directory(
       path_validation,
       seed=1993,
       image_size=image_size
   )
   

方案优势：

• 无需修改原有代码，零迁移成本；
• 新版本TensorFlow对GPU训练的性能优化更全面，API生态更完善；
• 官方镜像维护更可靠，避免手动构建的潜在问题。

方案二：在TF2.1.0中手动遍历目录加载数据 #

如果不想更换镜像，可通过遍历目录手动读取每张图片，转换为模型可用的数据集格式。

适配你的目录结构的实现代码：

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array

def load_dataset(dir_path, image_size=(90,90), color_mode='rgb'):
    """加载指定目录下的图像数据集"""
    X = []
    Y = []
    # 获取按名称排序的类别文件夹，保证类别索引一致
    class_folders = sorted([f.path for f in os.scandir(dir_path) if f.is_dir()])
    
    for class_idx, folder in enumerate(class_folders):
        # 遍历当前类别下的所有图片文件
        for img_file in os.listdir(folder):
            img_path = os.path.join(folder, img_file)
            # 跳过非图片文件（可选）
            if not img_file.lower().endswith(('.jpg', '.jpeg', '.png')):
                continue
            # 加载并预处理图片
            img = load_img(img_path, target_size=image_size, color_mode=color_mode)
            img_array = img_to_array(img).astype('float32') / 255.0  # 归一化到0-1区间
            X.append(img_array)
            Y.append(class_idx)
    
    # 转换为numpy数组
    X = np.array(X)
    Y = np.array(Y)
    # 获取类别名称
    class_names = [os.path.basename(f) for f in class_folders]
    return X, Y, class_names

# 加载训练集和验证集
train_X, train_Y, class_names = load_dataset(path_training, image_size=image_size)
val_X, val_Y, _ = load_dataset(path_validation, image_size=image_size)

# 转换为TensorFlow Dataset，支持批量加载、预取等优化
train_ds = tf.data.Dataset.from_tensor_slices((train_X, train_Y)) \
                          .shuffle(buffer_size=1000) \
                          .batch(32) \
                          .prefetch(tf.data.AUTOTUNE)

val_ds = tf.data.Dataset.from_tensor_slices((val_X, val_Y)) \
                        .batch(32) \
                        .prefetch(tf.data.AUTOTUNE)

方案优势与不足：

• 优势：无需更换镜像，可灵活自定义预处理逻辑；
• 不足：代码量较大，需自行处理数据加载的细节；数据量较大时，手动加载的效率不如官方API的流式加载机制。

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