加载Caltech TensorFlow数据集时遇InvalidArgumentError错误求助

这个错误是Caltech101数据集的常见问题——数据集中的图片原始尺寸并不统一（比如你遇到的[204,300,3]和[153,300,3]），而TensorFlow的batch()操作要求同一个batch内的所有张量必须具有完全相同的形状，因此才会抛出InvalidArgumentError。

直接解决方案：统一图片尺寸 #

在对数据集进行batch()操作之前，我们需要添加一个预处理步骤，把所有图片调整到相同的尺寸。这里提供两种常用的处理方式：

方式1：直接拉伸/缩放图片（简单高效）

使用tf.image.resize将所有图片统一缩放到指定尺寸（比如常用的224x224），代码修改如下：

import tensorflow as tf
import tensorflow_datasets as tfds

# 加载数据集
ds_train, ds_test = tfds.load(name="caltech101", split=["train", "test"])

# 定义预处理函数
def preprocess(features):
    # 统一图片尺寸为224x224，可根据你的模型需求调整
    resized_image = tf.image.resize(features["image"], (224, 224))
    # 可选：将像素值归一化到[0,1]，提升模型训练稳定性
    normalized_image = tf.cast(resized_image, tf.float32) / 255.0
    return {"image": normalized_image, "label": features["label"]}

# 先应用预处理，再执行shuffle、batch和prefetch
ds_train = ds_train.map(preprocess).shuffle(1000).batch(128).prefetch(10)

# 测试运行
for features in ds_train.take(1):
    image, label = features["image"], features["label"]
    print(f"Batch内图片形状：{image.shape}")
    print(f"Batch内标签形状：{label.shape}")

方式2：保持比例填充（避免图片变形）

如果你不想让图片被拉伸变形，可以使用tf.image.resize_with_pad，它会在图片的短边填充空白（默认补0），保持原始宽高比：

def preprocess(features):
    # 保持比例填充到224x224
    padded_image = tf.image.resize_with_pad(features["image"], 224, 224)
    normalized_image = tf.cast(padded_image, tf.float32) / 255.0
    return {"image": normalized_image, "label": features["label"]}

为什么这样能解决问题？ #

通过map()操作对每个样本单独处理后，所有图片的形状都会变成你指定的尺寸（比如(224,224,3)），此时再执行batch()操作，TensorFlow就能顺利将多个样本打包成一个形状统一的batch张量了。

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