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

结果:

参考:
https://docs.python.org/3/library/profile.html