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

Question

我正在使用python代码在FastAPI上部署预先训练好的机器学习模型。我正在使用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)

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

Answer 1

我对基于transformer的NLP模型的经验是，它们要求很高，有时会泄漏内存

例如，我对Facebook的Bart大型CNN模式（用于总结）存在问题，该模式的表现几乎相同（内存不断增加）

一个糟糕的解决方法是限制Docker容器使用的最大内存量，这样，一旦达到限制，Docker就会重新启动容器并释放所有内存。这不是100%满意，但它的工作，你不会打破你的主机