时间戳

首页 > 解决方案 > 使用 Flask 和 Docker 运行 TensorFlow 时出现内存问题

python - 使用 Flask 和 Docker 运行 TensorFlow 时出现内存问题

问题描述

我已经训练了一个 Tensorflow 对象检测模型，并希望将它与 Flask 一起部署。我的问题是，每次新请求时，docker 容器使用的内存都会增加 ~100mb，成功执行后不会释放。所以经过几次请求后，我的容器是 OOM。下面是我的烧瓶应用程序代码的片段。

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

在我看来，我使用@tf.function. 没有这个注解，when detect_fn(image)is not a graph 内存没有问题（波动但总是小于500mb）。但是作为 tf graph 执行这个检测函数要快 2 倍，所以我不想就此辞职。有没有办法解决这个内存使用问题？

标签： pythondockertensorflowflask

解决方案

推荐阅读