I've trained a Tensorflow object detection model and want to deploy it with Flask. My problem is that with every new request the memory used by docker container rises by ~100mb, which is not freed after successful execution. So after few requests my container is OOM. Below is the fragment of my flask app code.
detection_model = model_builder.build(model_config=configs['model'], is_training=False) ckpt = tf.compat.v2.train.Checkpoint(model=detection_model) ckpt.restore(os.path.join('model', 'ckpt-3')).expect_partial() @tf.function def detect_fn(image): image, shapes = detection_model.preprocess(image) prediction_dict = detection_model.predict(image, shapes) detections = detection_model.postprocess(prediction_dict, shapes) return detections app = Flask(__name__) @app.route('/detect', methods=['POST']) def detect(): image_np = get_img(request) input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32) detections = detect_fn(input_tensor) return Response(json.dumps(detections))
In my opinion there is something wrong with the way I use @tf.function. Without this annotation, when detect_fn(image) is not a graph there is no problem with memory (it fluctuates but is always less than 500mb). But executing this detection function as tf graph is 2x faster, so I would like to not resign from this. Is there any way to solve this problem with memory usage?