如何将整数序列转.tfrecords并还原数据集？CSV转TFRecords读取问题

我明白你在把CSV转成TFRecords再读取时遇到了麻烦，结合你的数据格式（50个整数特征+0/1标签），我给你一套完整的可运行方案，帮你搞定写入和读取的全流程：

一、正确写入TFRecords的代码 #

你之前用tf.python_io.TFRecordWriter的思路是对的，但要注意正确构建TFRecord的Example结构，确保特征和标签的格式匹配。这里是优化后的写入代码：

import tensorflow as tf
import csv

def csv_to_tfrecords(csv_file_path, tfrecords_file_path):
    # 初始化TFRecordWriter
    with tf.io.TFRecordWriter(tfrecords_file_path) as writer:
        with open(csv_file_path, 'r', newline='') as csvfile:
            reader = csv.reader(csvfile)
            # 跳过表头行
            next(reader)
            for row in reader:
                # 处理行内的空格（比如你的示例里"5 , 19"这种带空格的分隔）
                row = [item.strip() for item in row]
                # 提取50个特征和1个标签，转成整数
                features = list(map(int, row[:50]))
                label = int(row[50])
                
                # 构建TFRecord的Feature结构
                feature_dict = {}
                # 处理每个特征：因为是整数，用FixedLenFeature
                for i in range(50):
                    feature_dict[f'Feature{i+1}'] = tf.train.Feature(
                        int64_list=tf.train.Int64List(value=[features[i]])
                    )
                # 处理标签
                feature_dict['Label'] = tf.train.Feature(
                    int64_list=tf.train.Int64List(value=[label])
                )
                
                # 构建Example并序列化
                example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
                writer.write(example.SerializeToString())

# 调用示例
csv_to_tfrecords("your_data.csv", "output.tfrecords")

写入时的关键注意点： #

• 跳过表头行：一定要用next(reader)跳过第一行的Feature1到Label的表头，避免把表头当成数据写入。
• 处理空格：你的示例数据里每个字段带空格（比如"5 , 19"），所以用item.strip()去掉前后空格，避免转整数时报错。
• Feature类型匹配：因为你的特征和标签都是整数，所以用Int64List来存储，对应TFRecord的FixedLenFeature类型。

二、读取TFRecords并生成可用数据集的代码 #

写入完成后，需要正确解析TFRecords文件，转换成TensorFlow可用的tf.data.Dataset格式：

def parse_tfrecord_fn(example):
    # 定义解析格式，要和写入时的Feature结构完全对应
    feature_description = {}
    for i in range(50):
        feature_description[f'Feature{i+1}'] = tf.io.FixedLenFeature([], tf.int64)
    feature_description['Label'] = tf.io.FixedLenFeature([], tf.int64)
    
    # 解析Example
    example = tf.io.parse_single_example(example, feature_description)
    
    # 把特征整理成一个张量（可选，方便后续模型输入），标签转成int32（如果需要）
    features = tf.stack([example[f'Feature{i+1}'] for i in range(50)], axis=0)
    features = tf.cast(features, tf.float32)  # 如果模型需要浮点型特征，可以转成float32
    label = tf.cast(example['Label'], tf.int32)
    
    return features, label

def load_tfrecords_dataset(tfrecords_file_path, batch_size=32, shuffle=True):
    # 构建数据集
    dataset = tf.data.TFRecordDataset(tfrecords_file_path)
    # 解析每个Example
    dataset = dataset.map(parse_tfrecord_fn, num_parallel_calls=tf.data.AUTOTUNE)
    
    # 打乱和分批（根据需求调整）
    if shuffle:
        dataset = dataset.shuffle(buffer_size=1000)
    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    
    return dataset

# 调用示例
train_dataset = load_tfrecords_dataset("output.tfrecords", batch_size=32)

# 测试读取
for batch_features, batch_labels in train_dataset.take(1):
    print("Batch features shape:", batch_features.shape)
    print("Batch labels shape:", batch_labels.shape)

读取时的关键注意点： #

• 解析格式完全匹配：feature_description的键和类型必须和写入时的feature_dict完全一致，否则会解析失败。
• 特征整理：把50个单独的特征张量堆叠成一个形状为(batch_size, 50)的张量，这样更符合模型输入的要求。
• 性能优化：使用num_parallel_calls=tf.data.AUTOTUNE和prefetch来加速数据加载，适合大规模数据集。

如果之前的问题是写入时格式错误导致读取失败，或者读取时解析不匹配，用上面的代码应该能解决你的问题。

内容的提问来源于stack exchange，提问作者J. Schaefer