使用多处理池加速python中的TFLite推理

Question

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

该文档提到了调整使用线程数的可能性。但是我无法在 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 版本：1.14.0 从这里开始（x86-64 Python 3.5 版本）

蟒蛇版本：3.5

Answer 1

我知道这个帖子是两年半前创建的。

为我，

import multiprocessing
tf.lite.Interpreter(modelfile, num_threads=multiprocessing.cpu_count())

效果很好。

Answer 2

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

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