Ускоренный вывод TFLite в python с многопроцессорным пулом - PullRequest
Новая
       55

Ускоренный вывод TFLite в python с многопроцессорным пулом

0 голосов
/ 07 ноября 2019

Я играл с tflite и заметил на моем многоядерном процессоре, что он не сильно нагружен во время вывода. Я устранил узкое место ввода-вывода, предварительно создав случайные входные данные с пустыми значениями (случайные матрицы напоминают изображения), но затем tflite все еще не использует весь потенциал ЦП.

В документации упоминается возможность настройки количества используемых потоков. Однако я не смог выяснить, как это сделать в 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 из здесь (версия Python 3.5 для x86-64)

версия для Python: 3,5

...