0
$\begingroup$

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?

$\endgroup$

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.