修复加载模型期间tensorflow2中的内存泄漏

Question

我是tensorflow-2的新手，在对象检测应用程序中遇到内存泄漏。我能够跟踪导致问题的部分。

class TensorflowObjectDetector:
def __init__(self, saved_model_dir: str):
def detection_engine():
model = tf.saved_model.load(saved_model_dir)
return model.signatures['serving_default']
self.detection_engine = detection_engine()

def detect_frame(self, batch_of_frames: np.ndarray):   
input_tensor = tf.convert_to_tensor(batch_of_frames)
detections = self.detection_engine(input_tensor) # memory leak is with this line.

我发现了一些类似问题的信息https://github.com/tensorflow/tensorflow/issues/32234，https://github.com/tensorflow/models/issues/5139

在当前的应用程序中，我使用的是使用ssd_mobilenet训练的模型。如果我使用像faster_rcnn这样的不同模型来训练模型，内存问题会得到解决吗。或者有什么我遗漏的吗？但在添加了这两行之后，内存问题并没有完全解决，但能够减缓内存泄漏。

tf.keras.backend.clear_session()
gc.collect()

Answer 1

发现问题。我们使用tensorflow1预训练模型ssd_mobilenet_v2_coco，同时使用tensorflow2，切换到ssd mobilenet v2 FPNLite 640x640tensorflow2模型修复了内存泄漏