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()
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?


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.