Question

我正在与tflite一起玩，并且在我的多核CPU上观察到在推理期间压力并不大。我通过事先使用numpy创建随机输入数据（类似于图像的随机矩阵）消除了IO瓶颈，但是tflite仍然没有充分利用CPU的潜力。

documentation提到了调整已用线程数的可能性。但是，我无法在Python API中找到如何做到这一点。但是由于我已经看到人们为不同的模型使用多个解释器实例，所以我认为一个人可以使用同一模型的多个实例并在不同的线程/进程上运行它们。我编写了以下简短脚本：

import numpy as np
import os, time
import tflite_runtime.interpreter as tflite
from multiprocessing import Pool


# global, but for each process the module is loaded, so only one global var per process
interpreter = None
input_details = None
output_details = None
def init_interpreter(model_path):
    global interpreter
    global input_details
    global output_details
    interpreter = tflite.Interpreter(model_path=model_path)
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    interpreter.allocate_tensors()
    print('done init')

def do_inference(img_idx, img):
    print('Processing image %d'%img_idx)
    print('interpreter: %r' % (hex(id(interpreter)),))
    print('input_details: %r' % (hex(id(input_details)),))
    print('output_details: %r' % (hex(id(output_details)),))

    tstart = time.time()

    img = np.stack([img]*3, axis=2) # replicates layer three time for RGB
    img = np.array([img]) # create batch dimension
    interpreter.set_tensor(input_details[0]['index'], img )
    interpreter.invoke()

    logit= interpreter.get_tensor(output_details[0]['index'])
    pred = np.argmax(logit, axis=1)[0]
    logit = list(logit[0])
    duration = time.time() - tstart 

    return logit, pred, duration

def main_par():
    optimized_graph_def_file = r'./optimized_graph.lite'

    # init model once to find out input dimensions
    interpreter_main = tflite.Interpreter(model_path=optimized_graph_def_file)
    input_details = interpreter_main.get_input_details()
    input_w, intput_h = tuple(input_details[0]['shape'][1:3])

    num_test_imgs=1000
    # pregenerate random images with values in [0,1]
    test_imgs = np.random.rand(num_test_imgs, input_w,intput_h).astype(input_details[0]['dtype'])

    scores = []
    predictions = []
    it_times = []

    tstart = time.time()
    with Pool(processes=4, initializer=init_interpreter, initargs=(optimized_graph_def_file,)) as pool:         # start 4 worker processes

        results = pool.starmap(do_inference, enumerate(test_imgs))
        scores, predictions, it_times = list(zip(*results))
    duration =time.time() - tstart

    print('Parent process time for %d images: %.2fs'%(num_test_imgs, duration))
    print('Inference time for %d images: %.2fs'%(num_test_imgs, sum(it_times)))
    print('mean time per image: %.3fs +- %.3f' % (np.mean(it_times), np.std(it_times)) )



if __name__ == '__main__':
    # main_seq()
    main_par()

但是对于每个进程，通过hex(id(interpreter))打印的解释器实例的内存地址是相同的。但是，输入/输出详细信息的存储器地址不同。因此，我想知道，即使我可以加速，这种方式是否可能会出错？如果是这样，如何使用TFLite和python实现并行推理？

tflite_runtime版本：here的1.14.0（x86-64 Python 3.5版本）

python版本：3.5

Answer 1

我没有设置初始化程序，而是使用以下代码加载模型，并在同一函数中进行推断以解决此问题。

with Pool(processes=multiprocessing.cpu_count()) as pool:
   results = pool.starmap(inference, enumerate(test_imgs))

使用多处理池在python中加快TFLite推理

1 个答案: