python机器学习代码中的内存泄漏问题答案

【问题标题】：Memory leak issue in python Machine learning codepython机器学习代码中的内存泄漏问题
【发布时间】：2021-03-09 11:36:27
【问题描述】：

我正在使用 python 代码在快速 API 上部署预训练的机器学习模型。我正在使用 docker 将我的模型容器化。 API 代码工作正常，但在每个 API 命中后，它并没有释放它使用的内存。每次命中后，内存不断增加。这是我正在使用的代码，

import uvicorn
from fastapi import FastAPI
from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import pickle

from transformers.pipelines import pipeline

class QAItem(BaseModel):
    context: str
    question:str

with open('model_pickle','rb') as file:
    nlp = pickle.load(file)

app=FastAPI()

@app.post('/api')
def qanda(data:QAItem):
    data=data.dict()
    item=nlp(question=data['question'],
             context=data['context'],
             handle_impossible_answer=True,
             max_answer_len=500)
    json_compatible_item_data=jsonable_encoder(item)
    return JSONResponse(json_compatible_item_data)


if __name__=='__main__':
    uvicorn.run(app,host='0.0.0.0', port=8080)

有人知道我的代码有什么问题吗？

【问题讨论】：

我看到两个选项： 1 - nlp(...) 函数存在一些问题； 2 - Docker 有问题。我建议将与 ML 模型相关的代码注释掉并检查内存泄漏。否则，查看 Dockerfile 可能会有所帮助。
你能提供一个更新 - 你解决了这个问题吗？提前致谢

标签： python machine-learning memory-leaks fastapi

【解决方案1】：

我对基于 Transformer 的 NLP 模型的经验是，它们要求非常高，有时还会泄漏内存。

例如，我遇到了 Facebook 的 Bart 大型 CNN 模型（用于总结）的问题，该模型的行为方式几乎相同（内存不断增加）。

一个肮脏的解决方法是限制 Docker 容器使用的最大内存量，这样一旦达到限制，Docker 就会重新启动容器并释放所有内存。这不是 100% 令人满意，但它有效，而且您不会破坏您的主机...

【讨论】：