利用 onnxruntime 庫同時推理多個模型的效率研究

LyleChen發表於2022-04-06
模型

1. 背景

需求：針對視訊形式的資料輸入，對每一幀影像，有多個神經網路模型需要進行推理並獲得預測結果。如何讓整個推理過程更加高效，嘗試了幾種不同的方案。

硬體：單顯示卡主機。

2. 方案

由於存在多個模型需要推理，但模型之間沒有相互依賴關係，因此很容易想到通過並行的方式來提高執行效率。

對比了如下幾種方案的結果，包括：

  1. 序列
  2. 執行緒
  3. 程式
  4. 協程

3. 實現

3.1 整體流程

配置了 4 個體量相近的模型。
為了遮蔽讀取和解碼的時間消耗對最終結果的影響，提前讀取視訊並準備輸入。
統計每個單獨模型執行推理的累積時間，以及整體的執行時間。

import asyncio
from time import time

def main():
    frames = load_video()
    weights = load_weights()

    print('序列：')
    one_by_one(weights, frames)
    print('多執行緒：')
    multit_thread(weights, frames)
    print('多程式：')
    multi_process(weights, frames)
    print('協程：')
    asyncio.run(coroutine(weights, frames))

3.2 序列

讀取到當前幀資料後，所有模型依次執行。

def one_by_one(weights, frames):
    sessions = [init_session(weight) for weight in weights]
    costs = [[] for _ in range(len(weights))]
    since_infer = time()
    for frame in frames:
        for session in sessions:
            since = time()
            _ = session.run('output', {'input': frame})
            cost = time() - since
            costs[idx].append(cost)
    print([sum(cost) for cost in costs])
    print("infer:", time() - since_infer)
    return

3.3 多執行緒

為每一個模型分配一個執行緒。

from threading import Thread

def multit_thread(weights, frames):
    sessions = [init_session(weight) for weight in weights]
    threads = []
    since_infer = time()
    for session in sessions:
        thread = Thread(target=run_session_thread, args=(session, frames))
        thread.start()
        threads.append(thread)
    for thread in threads:
        thread.join()
    print("infer:", time() - since_infer)
    return

def run_session_thread(session, frames):
    costs = []
    for frame in frames:
        since = time()
        _ = session.run('output', {'input': frame})
        costs.append(time() - since)
    print(sum(costs))
    return

3.4 多程式

為每一個模型分配一個程式。
由於 session 不能在程式間傳遞，因此需要在每個程式的內部單獨初始化。如果資料較多，這部分初始化的時間消耗基本可以忽略不急。

from multiprocessing import Manager, Process

def multi_process(weights, frames):
    inputs = Manager().list(frames)
    processes = []
    since_infer = time()
    for weight in weights:
        process = Process(target=run_session_process, args=(weight, inputs))
        process.start()
        processes.append(process)
    for process in processes:
        process.join()
    print("infer:", time() - since_infer)
    return

def run_session_process(weight, frames):
    session = init_session(weight)
    costs = []
    for frame in frames:
        since = time()
        _ = session.run('output', {'input': frame})
        costs.append(time() - since)
    print(sum(costs))
    return

3.5 協程

為每一個模型分配一個協程。

async def coroutine(weights, frames):
    sessions = [init_session(weight) for weight in weights]
    since_infer = time()
    tasks = [
        asyncio.create_task(run_session_coroutine(session, frames))
        for session in sessions
    ]
    for task in tasks:
        await task
    print("infer:", time() - since_all)
    return

async def run_session_coroutine(session, frames):
    costs = []
    for frame in frames:
        since = time()
        _ = session.run('output', {'input': frame})
        costs.append(time() - since)
    print(sum(costs))
    return

3.6 其他輔助函式

import cv2
import numpy as np
import onnxruntime as ort

def init_session(weight):
    provider = "CUDAExecutionProvider"
    session = ort.InferenceSession(weight, providers=[provider])
    return session

def load_video():
    # 為了減少讀視訊的時間，複製相同的圖片組成batch
    vcap = cv2.VideoCapture('path_to_video')
    count = 1000
    batch_size = 4
    frames = []
    for _ in range(count):
        _, frame = vcap.read()
        frame = cv2.resize(frame, (256, 256)).transpose((2, 0, 1))
        frame = np.stack([frame] * batch_size, axis=0)
        frames.append(frame.astype(np.float32))
    return frames

def load_weights():
    return ['path_to_weights_0',
            'path_to_weights_1',
            'path_to_weights_2',
            'path_to_weights_3',]

4. 結果及分析

4.1 執行結果

以batch_size=4共執行 1000 幀資料，推理結果如下：

方案 序列 執行緒 程式 協程
單模型累積時間/s 7.9/5.3/5.2/5.2 13.5/13.5/15.6/15.7 13.5/13.8/13.7/13.6 6.5/5.2/5.3/5.3
總時間/s 23.7 15.8 30.1 22.5
視訊記憶體佔用/MB 1280 1416 3375 1280
平均 GPU-Util 約 60% 約 85% 約 70% 約 55%
  • 在這個場景下，多執行緒是綜合效率最高的方式（時間最短、視訊記憶體佔用合理、GPU 利用率最高）；
  • 序列作為最基礎的方案，總時間就是每個模型執行時間之和；
  • 多程式的方式，單模型的累積時間與多執行緒類似，但是總時間有明顯增加，且極大增加了視訊記憶體佔用；
  • 用協程的方式，總結果看，與序列模式本質上是一樣的。

4.2 結果分析

4.2.1 關於執行緒方案

為什麼多執行緒相比序列可以提高執行效率？

  • 基本的判斷是，session.run()函式執行時，既有 CPU 執行的部分，又有 GPU 執行的部分；
  • 如果是序列方案，則 CPU 執行時，GPU 會等待，反之亦然；
  • 當換用多執行緒方案後，當一個執行緒從 CPU 執行切換到 GPU 執行後，會繼續執行另一個執行緒的 CPU 部分，並等待 GPU 返回結果。

4.2.2 關於程式方案

為什麼多程式反而降低了執行效率？

  • 基本的判斷是，整體執行的瓶頸並不在 CPU 的運算部分，而是在於 GPU 上模型前向推理的計算部分；
  • 因此，用多個程式並沒有充分利用系統資源，多個 CPU 核心會爭奪同一個 GPU 的計算資源，並增加了排程消耗。

4.2.3 關於協程方案

為什麼看起來協程與序列的效果一樣？

協程方案在執行過程中，從表現上來看：

  • 單個模型的累積時間是逐步print出來的，間隔大致等於每個模型的累積時間（而執行緒和程式方案中，幾乎是同時輸出 4 個模型的累積時間，說明是同時執行結束）；
  • 視訊記憶體佔用是逐步增加的，最後達到與序列方案一致。

可能的原因：

  • CPU 和 GPU 的任務切換，可能無法觸發協程的切換，導致最終的效果是，一個模型完成了所有資料的推理後，再進行下一個模型的推理。

使用協程的必要性：

  • 從執行緒改為協程，是為了進一步降低執行緒切換的消耗；
  • 在這個場景下，需要同時執行推理的模型數量一般不會太多，建立同樣數量的執行緒，系統資源的消耗是可控的；
  • 因此，沒有使用協程的必要性。

關於協程的使用，也是現學，有可能因為使用方法不當而得出以上的結論。如有錯誤，歡迎指正。

相關文章