python 使用 cProfile 做 profiling
最近开始看机器学习的项目, 于是开始看 Python 的代码. 把一个机器学习的模型发布上 prod 去预测结果, 发现生产环境里面 的性能很差: 本地 1s 能跑完的 API, 在生产环境需要 30 多毫秒. 先是看了下基本情况, 发现生产环境在预测那段代码, 竟然起了 50 多个 Python 线程. 于是怀疑生产环境因为使用 container, 但是却拿到了宿主机的 CPU 数量, 于是开了很多线程. 但是 container 却限制了 cpu 的使用量, 导致多线程竞争, 最终性能下降.
于是尝试做 profiling: cProfile 是python 自带的.
要做 profiling 的部分:
import os
import time
import cProfile
from transformers import BertTokenizer, BertModel
pretrained_model_path = os.path.abspath(os.path.dirname(__file__)) + '/bert-base-uncased'
bert_tokenizer = BertTokenizer.from_pretrained(pretrained_model_path, cache_dir='/tmp')
bert_model = BertModel.from_pretrained(pretrained_model_path)
s = "This brings us to the downsides"
def bert_function():
t0 = time.time()
for i in range(0, 10):
inputs = bert_tokenizer(s, return_tensors="pt")
outputs = bert_model(**inputs)
print("used: {}".format((time.time() - t0)))
cProfile.run('bert_function()', 'my.prof')
执行
python test.py
使用 flameprof 转成 火焰图
python -m flameprof my.prof > my.svg
结果: