使用tf.py_function导致tf.data中ragged_batch失效的问题求助

针对你遇到的可变长度标注与tf.data流水线兼容的问题，以下是三种可行的解决方案：

方案1：修复tf.py_function的输出形状标注，使ragged_batch正常工作 #

问题核心：tf.py_function返回的张量默认没有明确的形状信息，导致ragged_batch无法识别可变长度的标注。手动设置每个张量的形状规则即可解决。

修改你的load函数，在tf.py_function调用后添加形状声明：

import tensorflow as tf
import xml.etree.ElementTree as ET

annotation_files = [
    '082f7a7f-IMG_0512.xml',
    '4f4c7f54-IMG_0511.xml',
    '5381454b-IMG_0510.xml',
    '05517884-IMG_0514.xml'
]
classNames = ["your_class1", "your_class2"]  # 替换为你的实际类别列表
imageFolder = "path/to/your/images"

def load(annotationFile):
    def _loadAnnotation(annotationFile):
        thisBoxes = []
        thisClassIDs = []
        annotationFile = annotationFile.numpy().decode("utf-8")
        root = ET.parse(annotationFile).getroot()
        for object in root.findall("object"):
            bndbox = object.find("bndbox")
            xmin = int(bndbox.find("xmin").text)
            ymin = int(bndbox.find("ymin").text)
            xmax = int(bndbox.find("xmax").text)
            ymax = int(bndbox.find("ymax").text)
            thisBoxes.append([xmin, ymin, xmax, ymax])
            className = object.find("name").text
            classID = classNames.index(className)
            thisClassIDs.append(classID)
        imageFile = imageFolder + "/" + root.find('filename').text
        return (imageFile, tf.cast(thisBoxes, dtype=tf.float32), tf.cast(thisClassIDs, dtype=tf.float32))
    
    imageFile, thisBoxes, thisClassIDs = tf.py_function(_loadAnnotation, [annotationFile], [tf.string, tf.float32, tf.float32])
    
    # 关键：手动设置张量形状，None表示维度长度可变
    imageFile.set_shape(())
    thisBoxes.set_shape((None, 4))
    thisClassIDs.set_shape((None,))
    
    # 加载并处理图像
    image = tf.io.read_file(imageFile)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.cast(image, tf.float32)
    image.set_shape((None, None, 3))
    
    bounding_boxes = {
        "boxes": thisBoxes,
        "classes": thisClassIDs
    }

    return {"images": image, "bounding_boxes": bounding_boxes}

dataset = tf.data.Dataset.from_tensor_slices(annotation_files)
dataset = dataset.map(load, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.ragged_batch(4)

核心逻辑：通过set_shape明确告诉TensorFlow每个张量的维度规则，让ragged_batch能正确识别可变长度的标注数据。

方案2：正确序列化TFRecord处理可变长度数据 #

TFRecord不支持嵌套Feature结构，且可变长度张量需要转换为扁平列表存储。同时用字节形式存储图像更高效。

序列化函数 #

def serializeTFRecord(data):
    # 将图像编码为JPEG字节，大幅减少磁盘占用
    image_bytes = tf.io.encode_jpeg(tf.cast(data["images"], tf.uint8)).numpy()
    boxes = data["bounding_boxes"]["boxes"].numpy()
    classes = data["bounding_boxes"]["classes"].numpy()
    num_boxes = boxes.shape[0]
    
    # 将边界框扁平化为一维列表
    boxes_flat = boxes.flatten()
    
    feature = {
        "image_bytes": tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_bytes])),
        "num_boxes": tf.train.Feature(int64_list=tf.train.Int64List(value=[num_boxes])),
        "boxes": tf.train.Feature(float_list=tf.train.FloatList(value=boxes_flat)),
        "classes": tf.train.Feature(float_list=tf.train.FloatList(value=classes))
    }
    
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

解析函数 #

def parse_tfrecord(example_proto):
    feature_description = {
        "image_bytes": tf.io.FixedLenFeature([], tf.string),
        "num_boxes": tf.io.FixedLenFeature([], tf.int64),
        "boxes": tf.io.VarLenFeature(tf.float32),
        "classes": tf.io.VarLenFeature(tf.float32)
    }
    parsed_features = tf.io.parse_single_example(example_proto, feature_description)
    
    # 解析图像
    image = tf.io.decode_jpeg(parsed_features["image_bytes"], channels=3)
    image = tf.cast(image, tf.float32)
    
    # 恢复边界框和类别ID的原始形状
    boxes = tf.sparse.to_dense(parsed_features["boxes"])
    boxes = tf.reshape(boxes, [parsed_features["num_boxes"], 4])
    classes = tf.sparse.to_dense(parsed_features["classes"])
    
    bounding_boxes = {
        "boxes": boxes,
        "classes": classes
    }
    return {"images": image, "bounding_boxes": bounding_boxes}

使用方式 #

# 先创建原始数据集（基于方案1的load函数）
raw_dataset = tf.data.Dataset.from_tensor_slices(annotation_files).map(load)

# 序列化并写入TFRecord
with tf.io.TFRecordWriter("detection_data.tfrecord") as writer:
    for data in raw_dataset:
        serialized_data = serializeTFRecord(data)
        writer.write(serialized_data)

# 加载并解析TFRecord
dataset = tf.data.TFRecordDataset("detection_data.tfrecord")
dataset = dataset.map(parse_tfrecord, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.ragged_batch(4)

优点：适合大规模数据集，磁盘占用低，加载速度快，完全兼容ragged_batch。

方案3：使用tf.data.Dataset.from_generator直接生成数据 #

如果数据集规模不大，可以直接用Python生成器解析XML，同时通过output_signature声明输出的RaggedTensor结构：

def data_generator(annotation_files):
    for ann_file in annotation_files:
        root = ET.parse(ann_file).getroot()
        thisBoxes = []
        thisClassIDs = []
        for obj in root.findall("object"):
            bndbox = obj.find("bndbox")
            xmin = int(bndbox.find("xmin").text)
            ymin = int(bndbox.find("ymin").text)
            xmax = int(bndbox.find("xmax").text)
            ymax = int(bndbox.find("ymax").text)
            thisBoxes.append([xmin, ymin, xmax, ymax])
            className = obj.find("name").text
            classID = classNames.index(className)
            thisClassIDs.append(classID)
        
        imageFile = imageFolder + "/" + root.find('filename').text
        image = tf.io.read_file(imageFile)
        image = tf.image.decode_jpeg(image, channels=3)
        image = tf.cast(image, tf.float32)
        
        yield {
            "images": image,
            "bounding_boxes": {
                "boxes": tf.constant(thisBoxes, dtype=tf.float32),
                "classes": tf.constant(thisClassIDs, dtype=tf.float32)
            }
        }

dataset = tf.data.Dataset.from_generator(
    lambda: data_generator(annotation_files),
    output_signature={
        "images": tf.TensorSpec(shape=(None, None, 3), dtype=tf.float32),
        "bounding_boxes": {
            "boxes": tf.RaggedTensorSpec(shape=(None, 4), dtype=tf.float32),
            "classes": tf.RaggedTensorSpec(shape=(None,), dtype=tf.float32)
        }
    }
)
dataset = dataset.ragged_batch(4)

优点：代码简洁，不需要处理tf.py_function的形状问题；缺点：并行处理效率略低于前两种方案，适合中小规模数据集。

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