python - 在 docker 图像中使用 transformers 包时出现问题

Question

我正在使用转换器管道对来自 6 种不同语言的示例文本执行情感分析。我在本地 Jupyterhub 中测试了代码，它运行良好。但是当我将它包装在一个 flask 应用程序中并从中创建一个 docker 镜像时，执行卡在管道推理线上并且需要很长时间才能返回情绪分数。

• mac os catalina 10.15.7(无 GPU)
• Python 版本:3.8
• 变形金刚包:4.4.2
• 手电筒版本:1.6.0

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
classifier = pipeline('sentiment-analysis', model=model, tokenizer=tokenizer)
results = classifier(["We are very happy to show you the Transformers library.", "We hope you don't hate it."])
print([i['score'] for i in results])

上面的代码在 Jupyter notebook 中运行良好，它为我提供了预期的结果

[0.7495927810668945,0.2365245819091797]

现在，如果我使用 flask wrapper 创建一个 docker 镜像，它会卡在 results = classifier([input_data]) 行，并且执行将永远运行。

我的文件夹结构如下:

- src
    |-- app
         |--main.py
    |-- Dockerfile
    |-- requirements.txt

我使用下面的 Dockerfile 来创建图像

FROM tiangolo/uwsgi-nginx-flask:python3.8
COPY ./requirements.txt /requirements.txt
COPY ./app /app
WORKDIR /app
RUN pip install -r /requirements.txt
RUN echo "uwsgi_read_timeout 1200s;" > /etc/nginx/conf.d/custom_timeout.conf

而我的requirements.txt文件如下:

pandas==1.1.5
transformers==4.4.2
torch==1.6.0

我的 main.py 脚本如下所示:

from flask import Flask, json, request, jsonify
import traceback
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline


app = Flask(__name__)
app.config["JSON_SORT_KEYS"] = False

model_name = 'nlptown/bert-base-multilingual-uncased-sentiment'
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
nlp = pipeline('sentiment-analysis', model=model_path, tokenizer=model_path)

@app.route("/")
def hello():
    return "Model: Sentiment pipeline test"


@app.route("/predict", methods=['POST'])
def predict():

    json_request = request.get_json(silent=True)
    input_list = [i['text'] for i in json_request["input_data"]]
    
    results = nlp(input_list)         ##########  Getting stuck here
    for result in results:
        print(f"label: {result['label']}, with score: {round(result['score'], 4)}")
    score_list = [round(i['score'], 4) for i in results]
    
    return jsonify(score_list)

if __name__ == "__main__":
    app.run(host='0.0.0.0', debug=False, port=80)

我的输入载荷是这样的

{"input_data" : [{"text" : "We are very happy to show you the Transformers library."},
                 {"text" : "We hope you don't hate it."}]}

我尝试查看变形金刚的 github 问题，但找不到。即使使用 flask 开发服务器，我的执行也能正常工作，但当我包装它并创建一个 docker 图像时，它会永远运行。我不确定在创建 docker 镜像时是否遗漏了要包含的任何其他依赖项。

谢谢。

Answer 1

我遇到了类似的问题。似乎启动应用程序会以某种方式污染变形金刚模型的内存。可能与 Flask 如何进行线程处理有关，但不知道为什么。为我解决问题的是在不同的线程中执行导致问题的事情(加载模型)。

import threading


def preload_models():
    "LOAD MODELS"
    return 0

def start_app():

    app = create_app()
    register_handlers(app)

    preloading = threading.Thread(target=preload_models)
    preloading.start()
    preloading.join()

    return app

先在这里回复。如果这对您有帮助，我将非常高兴。