批量加载tfrecord文件到keras模型

import tensorflow as tf
from tensorflow import keras

def parse_tfrecord_fn(example):
    feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64),
    }
    example = tf.io.parse_single_example(example, feature_description)
    image = tf.io.decode_image(example['image'], channels=3)
    image = tf.image.resize(image, (224, 224))
    label = example['label']
    return image, label

def load_tfrecord_dataset(file_pattern, batch_size=32):
    files = tf.data.Dataset.list_files(file_pattern)
    dataset = files.interleave(tf.data.TFRecordDataset, cycle_length=4)
    dataset = dataset.map(parse_tfrecord_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.shuffle(1000)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset

train_dataset = load_tfrecord_dataset('train.tfrecord', batch_size=32)
validation_dataset = load_tfrecord_dataset('validation.tfrecord', batch_size=32)

model = keras.applications.ResNet50(weights=None, classes=10)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_dataset, validation_data=validation_dataset, epochs=10)