如何在使用estimator训练期间仅将批量数据加载到内存中？

import tensorflow as tf

def input_fn():
    # 加载数据集
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    # 对数据集进行预处理和增强
    dataset = dataset.map(...)
    # 批量处理数据
    dataset = dataset.batch(batch_size)
    # 返回特征和标签
    return {'features': dataset}, {'labels': dataset}

def model_fn(features, labels, mode):
    # 定义模型结构和计算图
    ...

    if mode == tf.estimator.ModeKeys.TRAIN:
        # 训练模式
        loss = ...
        train_op = ...
        return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
    elif mode == tf.estimator.ModeKeys.EVAL:
        # 评估模式
        loss = ...
        eval_metric_ops = ...
        return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=eval_metric_ops)

# 创建Estimator
estimator = tf.estimator.Estimator(model_fn=model_fn)

# 配置训练和评估参数
train_spec = tf.estimator.TrainSpec(input_fn=input_fn, max_steps=num_train_steps)
eval_spec = tf.estimator.EvalSpec(input_fn=input_fn, steps=num_eval_steps)

# 训练和评估
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)