Fixing memory leak in tensorflow2 during loading the model

Question

I am new to tensorflow-2 and I am experiencing a memory leak in my object detection application. I was able to track the section which causes the issue.

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.

I found some information about similar issues https://github.com/tensorflow/tensorflow/issues/32234, https://github.com/tensorflow/models/issues/5139

In the current application, I am using a model trained using ssd_mobilenet if I use a different model like faster_rcnn to train the model will the memory issue be fixed. or is there anything that I am missing? But after adding these 2 lines memory issue is not fully fixed but was able to slow down the memory leak.

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

Answer 1

found the issue. we were using tensorflow1 pre-trained model ssd_mobilenet_v2_coco while working with tensorflow2, switching to SSD MobileNet V2 FPNLite 640x640 tensorflow2 model fixed the memory leak