在加载模型时修复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

发现了问题。我们在与ssd_mobilenet_v2_coco合作时使用了tensorflow1预训练的模型tensorflow2，切换到SSD MobileNet V2 FPNLite 640x640 tensorflow2模型修复了内存泄漏。