TensorFlow 在循环中减速

Question

我有多个循环运行的算法。包含 tensorflow 的那些在多次迭代后确实会变慢。

每个文件列表大约有 10,000 个文件，具体取决于它使用的算法。我一次遍历文件列表一个文件，从每个文件创建一个数据框，在数据框上运行我的算法，然后将结果写入数据库。看起来像：

file_list = self.get_files()
for file in file_list:
   data = self.get_data(file.fileid)
   result = self.get_result(data)
   self.write_result

get_result 是针对不同算法的不同函数。他们通常需要 0 - 5 秒来计算每个文件的结果。

我正在使用一种算法，在循环开始时每秒处理 2 个文件，但在几百个文件之后，它会减慢到每个文件一分钟。检查代码它必须是 TF 是瓶颈，因为其余代码相对微不足道。

在get_result 中有以下行，我认为是罪魁祸首：

z = self.evaluate_risk(normalized_X)

def evaluate_risk(self, X):    

    with tf.device('/cpu:0'):
        with tf.Session() as sess:
            tf.saved_model.loader.load(sess, model.pb)
            graph = tf.get_default_graph()
            input_x = graph.get_tensor_by_name("input:0")
            risk = graph.get_tensor_by_name("risk:0")
            z = sess.run(risk, {input_x: X})
            sess.close()
            del sess
            del graph
    return z

鉴于我使用的是with，我不明白为什么这个函数会导致任何问题。我已经添加了sess.close()、del sess 和del graph，但我仍然遇到同样的问题。

每次我有一个新文件并到达 result 时，我都应该从头开始 tensorflow。我的循环变慢有什么明显的原因吗？我猜 tensorflow 的某些部分没有重置。

Answer 1

没有看到一个完整的例子很难说什么是最好的解决方案，但通常我只会加载一次模型（可能在它自己的图表中）并只创建一个会话，然后在evaluate_risk中使用它.这应该会显着减少每次调用的开销。你可以这样做：

def __init__(self):
    # ... init code
    self.graph = tf.Graph()  # Have the model live in its own graph
    with self.graph.as_default(), tf.device('/cpu:0'):
        self.session = tf.Session()
        tf.saved_model.loader.load(self.session, model_pb)
        self.input_x = self.graph.get_tensor_by_name("input:0")
        self.risk = self.graph.get_tensor_by_name("risk:0")

def __del__(self):
    # Ensure the session is closed when the object is deleted
    # (or do it in another method, or make the object work as a context manager, ...)
    self.session.close()

def evaluate_risk(self, X):
    return self.session.run(self.risk, {self.input_x: X})

编辑：在__del__ 方法中关闭会话可能是多余的，因为原则上当对象被删除时，它的会话也将被关闭，从而关闭。但是，它避免了有人在对象中获取对会话的引用（如obj_session = my_object.session）的潜在问题，这可能导致会话未按预期关闭。它还使会话预计何时关闭更加清晰。