自定义模型训练循环中出现InvalidArgumentError问题求助

错误核心是训练数据中存在非JPEG格式的图片，或部分JPEG文件损坏，导致DecodeJpeg_1节点无法解析文件。结合你提供的代码，给出以下针对性解决步骤：

1. 排查数据集里的无效文件 #

批量检查数据集目录，找出所有非标准图片格式（或无后缀）的文件：

import os

dataset_dir = "替换成你的数据集根路径"
invalid_files = []
valid_extensions = [".jpg", ".jpeg", ".png", ".gif", ".bmp"]

for root, _, files in os.walk(dataset_dir):
    for file in files:
        ext = os.path.splitext(file)[1].lower()
        if ext not in valid_extensions:
            invalid_files.append(os.path.join(root, file))

print("找到无效格式文件：")
for f in invalid_files:
    print(f)

找到后可直接删除这些文件，或单独处理格式转换。

2. 修改预处理函数兼容多格式 #

原代码仅调用tf.image.decode_jpeg，只能处理JPEG格式，改成自动识别所有常见图片格式：

IMG_SIZE = 100

def preprocess_images(file_path, img_size=IMG_SIZE):
    # 读取文件
    image = tf.io.read_file(file_path)
    # 自动解码所有支持的图片格式（JPEG/PNG/GIF/BMP）
    image = tf.io.decode_image(image, channels=3, expand_animations=False)
    # 归一化+调整尺寸
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = tf.image.resize(image, size=[img_size, img_size])
    return image

如果需要精准控制格式，也可按扩展名分支处理：

def preprocess_images(file_path, img_size=IMG_SIZE):
    image = tf.io.read_file(file_path)
    # 获取文件扩展名
    ext = tf.strings.split(file_path, ".")[-1]
    # 根据扩展名选择解码方式
    image = tf.cond(
        tf.logical_or(tf.equal(ext, "jpg"), tf.equal(ext, "jpeg")),
        lambda: tf.image.decode_jpeg(image, channels=3),
        lambda: tf.image.decode_png(image, channels=3)
        # 可继续添加GIF/BMP的解码逻辑
    )
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = tf.image.resize(image, size=[img_size, img_size])
    return image

3. 过滤损坏的JPEG文件 #

部分JPEG文件可能损坏，解码时触发错误，可在预处理中添加容错逻辑：

def preprocess_images(file_path, img_size=IMG_SIZE):
    def safe_decode(image_data):
        try:
            return tf.image.decode_jpeg(image_data, channels=3)
        except tf.errors.InvalidArgumentError:
            # 损坏文件返回空白图，也可返回None后续过滤
            return tf.zeros([img_size, img_size, 3], dtype=tf.uint8)

    image = tf.io.read_file(file_path)
    # 使用py_function包裹容错解码逻辑
    image = tf.py_function(safe_decode, [image], tf.uint8)
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = tf.image.resize(image, size=[img_size, img_size])
    return image

也可在构建Dataset时直接过滤无效样本：

def is_valid_image(file_path):
    try:
        image = tf.io.read_file(file_path)
        tf.image.decode_jpeg(image, channels=3)
        return True
    except tf.errors.InvalidArgumentError:
        return False

# 过滤数据集里的无效图片
train_data = train_data.filter(lambda x, y, z: is_valid_image(x))

4. 确认数据集完整性 #

• 检查最近是否新增过不符合格式的图片到数据集
• 确认训练代码加载的数据集路径是否正确，是否误加载了其他目录的文件

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