使用tf.data API、TFRecordDataset及序列化时遇到问题

我懂你现在的痛点——手里的大型numpy整数数据集塞不下GPU内存，想靠TFRecord流式读取来解决，但原始数据的维度是非均匀的，直接用常规的均匀维度序列化方法肯定行不通。别急，我给你一套完整的解决方案，从数据生成、序列化到流式读取解析全给你捋清楚。

核心问题拆解 #

TFRecord默认是为固定形状的张量设计的，非均匀（变长）的数据需要用tf.VarLenFeature或者更灵活的tf.RaggedTensor来处理。后者在TensorFlow 2.x里支持得更好，咱们重点用它来实现。

步骤1：生成模拟非均匀numpy数据 #

先造一批和你场景类似的非均匀整数数据，比如每个样本是长度不一的一维整数数组：

import numpy as np

# 生成100个非均匀长度的整数数组
non_uniform_data = []
for _ in range(100):
    # 每个数组长度随机在5到20之间
    seq_len = np.random.randint(5, 21)
    seq = np.random.randint(0, 100, size=seq_len, dtype=np.int64)
    non_uniform_data.append(seq)

步骤2：将非均匀数据序列化到TFRecord #

对于变长整数序列，咱们用tf.train.Int64List来存储，这样比转成bytes更高效，也更贴合整数数据的特性：

import tensorflow as tf

def serialize_non_uniform_data(seq):
    # 把numpy数组转成tf.train.Int64List
    feature = {
        'sequence': tf.train.Feature(int64_list=tf.train.Int64List(value=seq))
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

# 写入TFRecord文件
with tf.io.TFRecordWriter('non_uniform_data.tfrecord') as writer:
    for seq in non_uniform_data:
        serialized_example = serialize_non_uniform_data(seq)
        writer.write(serialized_example)

如果你的数据是更高维的非均匀数据（比如每个样本是形状不一的二维数组），那需要额外存储形状信息，再把数组转成序列化格式：

def serialize_2d_non_uniform_data(arr):
    feature = {
        'shape': tf.train.Feature(int64_list=tf.train.Int64List(value=arr.shape)),
        'data': tf.train.Feature(int64_list=tf.train.Int64List(value=arr.flatten()))
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

步骤3：流式读取并解析TFRecord #

读取时，咱们用tf.io.parse_single_example解析，并把变长特征转成tf.RaggedTensor，方便后续灵活处理：

def parse_example(serialized_example):
    # 定义特征描述：sequence是变长int64
    feature_descriptions = {
        'sequence': tf.io.VarLenFeature(tf.int64)
    }
    parsed_features = tf.io.parse_single_example(serialized_example, feature_descriptions)
    # 将VarLenFeature转成RaggedTensor，支持后续灵活操作
    sequence = tf.RaggedTensor.from_sparse(parsed_features['sequence'])
    return sequence

# 构建数据集流式读取管道
dataset = tf.data.TFRecordDataset('non_uniform_data.tfrecord')
dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)

# 测试读取效果
for seq in dataset.take(3):
    print(f"序列长度：{seq.shape[0]}，内容：{seq.numpy()}")

如果是解析二维非均匀数据，解析函数要改成这样：

def parse_2d_example(serialized_example):
    feature_descriptions = {
        'shape': tf.io.FixedLenFeature([2], tf.int64),
        'data': tf.io.VarLenFeature(tf.int64)
    }
    parsed_features = tf.io.parse_single_example(serialized_example, feature_descriptions)
    # 还原形状和数组
    shape = parsed_features['shape']
    data = tf.sparse.to_dense(parsed_features['data'])
    arr = tf.reshape(data, shape)
    return arr

后续处理注意事项 #

• 如果要把数据喂给GPU模型，非均匀数据通常需要做padding处理，比如用tf.data.Dataset.padded_batch把一个batch里的序列补到相同长度：

  dataset = dataset.padded_batch(batch_size=8, padded_shapes=([None]))
  

• TensorFlow的很多层（比如tf.keras.layers.Embedding）已经支持RaggedTensor，可以直接输入，不需要额外padding。
• 流式读取时，设置num_parallel_calls=tf.data.AUTOTUNE能自动利用空闲CPU资源加速map处理，提升整体读取效率。

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