linux中tfrecords

import tensorflow as tf

# 定义特征描述
feature_description = {
    'image': tf.io.FixedLenFeature([], tf.string),
    'label': tf.io.FixedLenFeature([], tf.int64),
}

def serialize_example(image, label):
    feature = {
        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),
        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

# 写入 TFRecords 文件
with tf.io.TFRecordWriter('data.tfrecords') as writer:
    for image, label in dataset:
        example = serialize_example(image.numpy(), label.numpy())
        writer.write(example)

def parse_example(serialized_example):
    features = tf.io.parse_single_example(serialized_example, feature_description)
    image = tf.io.decode_raw(features['image'], tf.uint8)
    label = tf.cast(features['label'], tf.int32)
    return image, label

dataset = tf.data.TFRecordDataset(['data.tfrecords'])
dataset = dataset.map(parse_example)