使用Dask分块大型图像数据集后，CNN训练过慢如何解决？

我正在处理一个形状为(10000000,1,32,32)（样本数、通道数、高、宽）的大型图像数据集，已用Dask完成数据分块，代码如下：

import dask.array as da
X = da.from_array(f['features'], chunks=(1000, 1, 32,32))
y =  da.from_array(f['targets'], chunks=(1000,1))

我的GPU内存仅6GB，想借助Dask分块训练CNN，但运行时出现以下警告：

WARNING:tensorflow:Keras is training/fitting/evaluating on array-like data. Keras may not be optimized for this format, so if your input data format is supported by TensorFlow I/O (https://github.com/tensorflow/io) we recommend using that to load a Dataset instead.

且单轮epoch训练耗时约30小时。作为Dask新手，查阅文档后仍感困惑，恳请提供解决方法。

CNN模型代码如下：

lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, min_lr=1.e-6)
early_stopping_cb = tf.keras.callbacks.EarlyStopping(monitor="val_loss",patience=5, restore_best_weights=True)
callbacks = [lr_scheduler, early_stopping_cb]

inputs1 = keras.Input(shape=(1,32,32), batch_size = 64)
x1 = Conv2D(32, kernel_size=2,  padding='same', strides = 1, activation='relu', kernel_initializer='TruncatedNormal',
           data_format='channels_first')(inputs1)
x1 = Conv2D(32,  kernel_size=2, padding='same', strides = 1, activation='relu',kernel_initializer='TruncatedNormal',
           data_format='channels_first')(x1)
x1 = MaxPooling2D(pool_size=(2, 2), data_format='channels_first')(x1)
x1 = Conv2D(64,  kernel_size=3, padding='same', strides = 1, activation='relu',kernel_initializer='TruncatedNormal',
           data_format='channels_first')(x1)
x1 = Conv2D(64,  kernel_size=3, padding='same', strides = 1, activation='relu',kernel_initializer='TruncatedNormal',
           data_format='channels_first')(x1)
x1 = Flatten()(x1)
x1 = Dense(128, activation='relu', kernel_initializer='TruncatedNormal')(x1)
x1 = Dropout(0.2)(x1)
x1 = Dense(32, activation='relu', kernel_initializer='TruncatedNormal')(x1)
x1 = Dropout(0.2)(x1)
outputs1 = Dense(1, activation='sigmoid', kernel_initializer='TruncatedNormal')(x1)
conv2d = keras.Model(inputs=inputs1, outputs=outputs1)
conv2d.compile(loss = "binary_crossentropy", optimizer = "adam", metrics = ["accuracy"])
conv2d.fit(X,y, batch_size=64, epochs = 30,
                    callbacks=callbacks)

1. 将Dask数组转为TensorFlow Dataset适配Keras #

Keras直接处理Dask数组效率低下，因为没有针对性优化。可以把Dask分块数据转为tf.data.Dataset，利用TF的预取、并行加载能力提升速度：

import tensorflow as tf
import dask.array as da

def dask_to_tf_dataset(X_dask, y_dask, batch_size=64):
    def generator():
        # 逐个加载Dask块，再按batch拆分
        for x_chunk, y_chunk in zip(X_dask.to_delayed(), y_dask.to_delayed()):
            x_np = x_chunk.compute()
            y_np = y_chunk.compute()
            for i in range(0, len(x_np), batch_size):
                yield x_np[i:i+batch_size], y_np[i:i+batch_size]
    
    input_shape = X_dask.shape[1:]
    output_shape = y_dask.shape[1:]
    
    return tf.data.Dataset.from_generator(
        generator,
        output_signature=(
            tf.TensorSpec(shape=(None,) + input_shape, dtype=tf.float32),
            tf.TensorSpec(shape=(None,) + output_shape, dtype=tf.float32)
        )
    ).prefetch(tf.data.AUTOTUNE)

# 转换数据集
train_dataset = dask_to_tf_dataset(X, y, batch_size=64)
# 训练时改用Dataset输入
conv2d.fit(train_dataset, epochs=30, callbacks=callbacks)

2. 调整Dask分块大小 #

当前(1000,1,32,32)的分块过小，会导致频繁的调度开销。建议将分块大小设为batch_size的整数倍，比如(640,1,32,32)（64*10），减少调度次数：

X = da.from_array(f['features'], chunks=(640, 1, 32,32))
y = da.from_array(f['targets'], chunks=(640,1))

3. 启用Dask GPU加速（可选） #

如果环境支持，用dask-cuda让Dask直接在GPU上处理数据，减少CPU-GPU数据传输耗时：

from dask_cuda import LocalCUDACluster
from dask.distributed import Client

# 创建本地GPU集群
cluster = LocalCUDACluster()
client = Client(cluster)

# 将Dask数组转为GPU张量
X = X.map_blocks(lambda x: tf.convert_to_tensor(x, dtype=tf.float32))
y = y.map_blocks(lambda x: tf.convert_to_tensor(x, dtype=tf.float32))

4. 模型训练细节优化 #

• 移除Input层的batch_size参数，让TF自动适配Dataset的batch大小
• 启用混合精度训练，降低GPU内存占用并提升速度：

from tensorflow.keras import mixed_precision
mixed_precision.set_global_policy('mixed_float16')
# 注意输出层保持float32避免精度损失
outputs1 = Dense(1, activation='sigmoid', kernel_initializer='TruncatedNormal', dtype='float32')(x1)

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