我正在尝试在 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
...
...