时间戳

首页 > 解决方案 > 首次运行 TorchScript 模型及其在 Flask 中的实现太慢

python - 首次运行 TorchScript 模型及其在 Flask 中的实现太慢

问题描述

我正在尝试在 Python 和 Flask 中部署 torchscripted 模型。正如我意识到（至少如此处所述）脚本模型在使用前需要“预热”，因此此类模型的第一次运行需要比后续模型更长的时间。我的问题是：有没有办法在 Flask 路由中加载torchscripted 模型并在不损失“蠕虫”时间的情况下进行预测？我可以将“预热”模型存储在某个地方以避免在每个请求中进行预热吗？我编写了简单的代码来重现“热身”过程：

    import torchvision, torch, time
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
    model = torch.jit.script(model)
    model.eval()
    x = [torch.randn((3,224,224))]
    for i in range(3):
      start = time.time()
      model(x)
      print(‘Time elapsed: {}’.format(time.time()-start))

输出：

Time elapsed: 38.29<br>
Time elapsed: 6.65<br>
Time elapsed: 6.65<br>

和烧瓶代码：

    import torch, torchvision, os, time
    from flask import  Flask
    app = Flask(__name__)
    @app.route('/')
    def test_scripted_model(path='/tmp/scripted_model.pth'):
        if os.path.exists(path):
            model = torch.jit.load(path, map_location='cpu')
        else:
            model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
            model = torch.jit.script(model)
            torch.jit.save(model, path)
        model.eval()
        x = [torch.randn((3, 224, 224))]
        out = ''
        for i in range(3):
            start = time.time()
            model(x)
            out += 'Run {} time: {};\t'.format(i+1, round((time.time() - start), 2))
        return out
    if __name__ == '__main__':
        app.run(host='0.0.0.0', port=5000, debug=False)

输出：

Run 1 time: 46.01; Run 2 time: 8.76; Run 3 time: 8.55;

操作系统：Ubuntu 18.04 和 Windows10
Python 版本：3.6.9 Flask：1.1.1 Torch：1.4.0 Torchvision：0.5.0

更新：

将“热身”问题解决为：

with torch.jit.optimized_execution(False):
    model(x)

Update2： 解决了 Flask 问题（如下所述），在服务器启动之前创建全局 python 模型对象并在那里对其进行预热。然后在每个请求中，模型都可以使用了。

    model = torch.jit.load(path, map_location='cpu').eval()
    model(x)
    app = Flask(__name__)

然后在@app.route 中：

@app.route('/')
def test_scripted_model():
    global model
    ...
    ...

标签： pythonflaskpytorchtorchscript

解决方案

我可以将“预热”模型存储在某个地方以避免在每个请求中进行预热吗？

是的，只需在函数外部实例化您的模型test_scripted_model并从函数内部引用它。

推荐阅读