在Tensorflow 2中发布keras会话资源

如何解决在Tensorflow 2中发布keras会话资源

我在Tensorflow 1.x上运行了一个keras模型，在其中定义了一些数据集，然后将其应用于keras模型。代码结构如下：

import tensorflow as tf
from tensorflow.keras import backend as K
config = tf.ConfigProto()
config.gpu_options.allow_growth = True

def doWork(args):

    sess = tf.Session(config=config)
    K.set_session(sess)

    ... generate datasets,create keras model,run training loop,etc.

    K.clear_session()
    sess.close()

if __name__ == '__main__':
    
    ... set some global variables
    doWork(first set of arguments)
    doWork(second set of arguments)
    ...

我想在开始时（configProto行）限制gpu内存使用，因此我需要使用所需的gpu配置定义会话。我也将数据集缓存在计算机主存储器中以加快训练过程，并希望在完成训练后将其发布。因此，我将训练循环放入doWork（）函数中，以便能够在每次运行它时分配和释放资源。我发现除非关闭会话，否则不会释放内存。我还发现了一个与keras的clear_session并添加了它（但是我不确定它是否对我有帮助）。

上面的代码在TF 1.x上可以正常工作，但是我不知道如何将其移植到TF 2，因为会话已从那里的常规使用中删除。如何做到这一点，并在TF 2中完成configProto和主内存释放？我应该再次定义会话吗？我认为应该有更简单的方法。