Colab使用TPU训练CNN时ImageDataGenerator数据增强报错解决方案问询

1 先修复TPU初始化代码的显性疏漏 #

你现有TPU初始化代码中创建TPUStrategy时未传入resolver参数，属于可直接触发报错的bug，先修正：

# 原有错误写法：tpu_strategy = tf.distribute.TPUStrategy()
# 修改为
tpu_strategy = tf.distribute.TPUStrategy(resolver)

同时注意：AlexNet模型的定义、编译代码必须全部放在tpu_strategy.scope()上下文管理器中，否则TPU无法正确分发模型参数，示例如下：

with tpu_strategy.scope():
    # 此处放入你所有的AlexNet模型定义、compile代码
    alexnet_cnn = Sequential()
    # ...... 剩余模型结构代码
    alexnet_cnn.compile(optimizer = 'Adam', loss = 'categorical_crossentropy', metrics = ['accuracy'])

2 最优适配方案：替换ImageDataGenerator为tf.data.Dataset实现相同增强逻辑 #

TPU原生不支持Python生成器格式输入（ImageDataGenerator返回的DataFrameIterator本质就是Python生成器），必须改用tf.data格式输入，以下是和你现有增强参数完全对齐的实现代码：

import tensorflow as tf

image_size = 224
batch_size = 16
# 你的任务为4分类
NUM_CLASSES = 4

# 单张图像增强函数，和你原有ImageDataGenerator配置完全对应
def augment_image(image, label):
    # 对应rescale=1/255
    image = tf.cast(image, tf.float32) / 255.0
    # 对应rotation_range=15
    image = tf.keras.layers.RandomRotation(factor=15/360, fill_mode='nearest')(image)
    # 对应width_shift_range=0.1
    image = tf.keras.layers.RandomTranslation(height_factor=0, width_factor=0.1, fill_mode='nearest')(image)
    # 对应zoom_range=0.01
    image = tf.keras.layers.RandomZoom(height_factor=0.01, width_factor=0.01, fill_mode='nearest')(image)
    # 对应shear_range=0.01
    image = tf.keras.layers.RandomShear(x_factor=0.01, y_factor=0.01, fill_mode='nearest')(image)
    # 对应brightness_range=[0.3, 1.5]
    image = tf.keras.layers.RandomBrightness(factor=0.5, value_range=(0,1))(image)
    image = tf.clip_by_value(image, 0.0, 1.0)
    # 对应horizontal_flip、vertical_flip均为True
    image = tf.keras.layers.RandomFlip(mode="horizontal_and_vertical")(image)
    return image, label

# 从你原始DataFrame生成tf.data格式数据集
def create_dataset(df, is_training=True):
    file_paths = df['image_filepaths'].values
    labels = tf.keras.utils.to_categorical(df['tumor_class'], num_classes=NUM_CLASSES)
    
    ds = tf.data.Dataset.from_tensor_slices((file_paths, labels))
    
    # 加载图像逻辑
    def load_image(file_path, label):
        img = tf.io.read_file(file_path)
        # 你的图像为灰度图，通道数为1
        img = tf.io.decode_png(img, channels=1)
        img = tf.image.resize(img, (image_size, image_size))
        return img, label
    
    ds = ds.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
    
    # 训练集加数据增强，验证、测试集默认不加，可按需调整
    if is_training:
        ds = ds.map(augment_image, num_parallel_calls=tf.data.AUTOTUNE)
    
    # 分批、预取适配TPU性能优化
    ds = ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return ds

# 生成对应数据集
train_dataset = create_dataset(train_data, is_training=True)
validation_dataset = create_dataset(validation_data, is_training=False)
test_dataset = create_dataset(test_data, is_training=False)

修改完成后直接使用你原有fit代码即可正常在TPU上训练。

3 备选方案：将DataFrameIterator转换为TFRecord格式后读取 #

如果你不想修改现有数据增强逻辑，可以先把DataFrameIterator生成的样本存储为TPU支持的TFRecord格式，再加载训练，实现代码如下：

import tensorflow as tf

# 将迭代器内容写入TFRecord
def create_tfrecord_from_iterator(iterator, output_path, total_samples):
    writer = tf.io.TFRecordWriter(output_path)
    step_count = total_samples // iterator.batch_size
    for _ in range(step_count):
        images, labels = iterator.next()
        for i in range(images.shape[0]):
            # 图像转回uint8存储减少体积
            img_bytes = tf.io.encode_png(tf.cast(images[i]*255, tf.uint8)).numpy()
            label_bytes = labels[i].tobytes()
            feature = {
                'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_bytes])),
                'label': tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_bytes]))
            }
            example = tf.train.Example(features=tf.train.Features(feature=feature))
            writer.write(example.SerializeToString())
    writer.close()

# 分别生成训练、验证、测试集TFRecord
create_tfrecord_from_iterator(train_dataset, 'train.tfrecord', len(train_data))
create_tfrecord_from_iterator(validation_dataset, 'val.tfrecord', len(validation_data))
create_tfrecord_from_iterator(test_dataset, 'test.tfrecord', len(test_data))

# 解析TFRecord生成数据集
def parse_tfrecord_fn(example):
    feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.string),
    }
    example = tf.io.parse_single_example(example, feature_description)
    image = tf.io.decode_png(example['image'], channels=1)
    image = tf.cast(image, tf.float32) / 255.0
    label = tf.io.decode_raw(example['label'], tf.float32)
    label = tf.reshape(label, (4,))
    return image, label

# 加载TFRecord数据集
def load_tfrecord_dataset(file_path):
    ds = tf.data.TFRecordDataset(file_path)
    ds = ds.map(parse_tfrecord_fn, num_parallel_calls=tf.data.AUTOTUNE)
    ds = ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return ds

train_ds = load_tfrecord_dataset('train.tfrecord')
val_ds = load_tfrecord_dataset('val.tfrecord')
test_ds = load_tfrecord_dataset('test.tfrecord')

使用生成的TFRecord数据集即可正常在TPU上训练。

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