Python 如何为tf.data.Dataset构建外部缓存

我想知道如何写一个正确的输入法。我使用Tensorflow估计器和数据集重建了keras UNet实现。它正在工作，但比Keras实现慢得多。 当Keras实现将整个数据集加载到内存时，我的输入管道直接从SSD读取文件

我想在数据集中做同样的事情，不幸的是，

ds.cache（）

生成的缓存在每次中断

estimator.train

时都会被清除。我必须中断它来验证模型，因为我试图只使用一个GPU

我看了一下Dataset.cache的实现，我猜缓存绑定到保存数据集的图，每次训练中断时，这个图都会被释放

我有哪些选项可以使我的输入管线保持在tensorflow图中？我希望避免在python和tensorflow之间来回传递整个培训数据

我可以在另一个图中使用一个图中的张量吗？所以我可以用我的输入创建一个外部（到Estimator的API）图，然后在Estimator.train创建的图中使用它

大致如下：

train_graph = tf.Graph()
with training_graph.as_default():
    train_get_next = input_fn(training_set, ...) () # tensor returned by it.get_next()

def train_input_fn_with_ext_cache():
    batch = # some how obtain result of train_get_next 
    return batch

train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn_with_ext_cache,  max_steps=3000)
eval_spec = ...

tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)