使用load()方法从Tensorflow数据集中加载数据集的一部分

import tensorflow as tf

dataset_path = "path/to/dataset"
dataset_file = "dataset.tfrecord"

feature_description = {
    'feature1': tf.io.FixedLenFeature([], tf.int64),
    'feature2': tf.io.FixedLenFeature([], tf.float32),
    'feature3': tf.io.FixedLenFeature([], tf.string),
}

def parse_example(example_proto):
    return tf.io.parse_single_example(example_proto, feature_description)

def load_dataset(dataset_path, dataset_file):
    dataset = tf.data.TFRecordDataset(dataset_path + "/" + dataset_file)
    dataset = dataset.map(parse_example)
    return dataset

dataset = load_dataset(dataset_path, dataset_file)