Python > Docker > DockerでGPU版ONNXを使ってみる
   こしあん
2022-12-04

DockerでGPU版ONNXを使ってみる

Posted On 2022-12-04

このエントリーをはてなブックマークに追加
Pocket
LINEで送る
Delicious にシェア

2.4k{icon} {views}


DockerでGPU版をONNXを動かしてみました。比較的に簡単に動かせたので、いろいろ便利だと思います。YOLOXを例に検証します。

概要

ONNXRuntimeのGPU版をDockerで使うことをテストします。YOLOXをサンプルに動かしてみます。

具体的にやることは以下の通りです。

  • 動画内の全フレームを物体検出（YOLOX）で推論する
  • YOLOXで公開されているONNXモデルを使う
  • 推論環境をCPU版とGPU版で変えて、Docker上で動かしてみる。推論速度を比較する

GPU版とCPU版の切り替え

docker run時にGPUを使うかどうかのオプションがあります。例えば、イメージ名がonnx_gpuとしたとき、

docker run --rm -it --gpus all onnx_gpu

とすればGPUあり

docker run --rm -it onnx_gpu

と「gpus」オプションを切れば、GPUなしとなります。

一方で、ONNX側の切り替えはどうすればよいでしょうか。セッションのProviderを指定することで切り替えができます。こちらの記事を参考にしました。

    try:
        provider = ['CUDAExecutionProvider','CPUExecutionProvider']
        session = ort.InferenceSession(f'yolox/{args.model}.onnx', providers=provider)
    except:
        provider = ['CPUExecutionProvider']
        session = ort.InferenceSession(f'yolox/{args.model}.onnx', providers=provider)
    print(session.get_providers())

Docker内でGPUが利用できるかどうかは、多分いろいろやり方あると思うのですが、雑にtry～exceptでセッションを起動しました。ちなみに、GPUが利用できない状態（gpusのオプションを切った状態で起動）でProviderを「’CUDAExecutionProvider’,’CPUExecutionProvider’」と指定すると、以下のようなエラーが出ます。

RuntimeError: /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:124 std::conditional_t<THRW, void, onnxruntime::common::Status> onnxruntime::CudaCall(ERRTYPE, const char*, const char*, ERRTYPE, const char) [with ERRTYPE = cudaError; bool THRW = true; std::conditional_t<THRW, void, onnxruntime::common::Status> = void] /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:117 std::conditional_t<THRW, void, onnxruntime::common::Status> onnxruntime::CudaCall(ERRTYPE, const char, const char*, ERRTYPE, const char*) [with ERRTYPE = cudaError; bool THRW = true; std::conditional_t<THRW, void, onnxruntime::common::Status> = void] CUDA failure 35: CUDA driver version is insufficient for CUDA runtime version ; GPU=0 ; hostname=05e86dc49193 ; expr=cudaSetDevice(info_.device_id);

このエラーをハンドリングして、セッションのProviderを再指定すればCPUで起動するというわけです。

ディレクトリ構成

  • src
    • main.py
    • yolox_utils.py
  • yolox
    • target_video.mp4
    • yolox_nano.onnx
    • yolox_tiny.onnx
    • yolox_s.onnx
    • yolox_m.onnx
    • yolox_l.onnx
    • yolox_x.onnx
  • Dockerfile

ONNXファイルは公式リポジトリからDLしました。詳細なソースは記事末尾を参照してください。

yolox_utils.pyはNMSや前処理のファイルです。リポジトリからコピペしてきました（このためだけにYOLOXとPyTorchをインストールするのもだるかったので）

target_video.mp4は推論したい動画です。こちらのフリー素材を使いました。実際はHD解像度の動画を使っています。

main.pyは、YOLOXで動画の全フレームに対して物体検出をかけます。単にONNXモデル推論するだけでなく、前処理＋NMSや後処理も含みます（Bounding Boxの可視化はしません）。動画の読み込みはOpenCVを使いました。

ハードウェア

  • GPU: RTX 2080Ti
  • CPU: Core i9-9900K @ 3.60GHz

Dockerのビルドと実行

Dockerfileをビルドします

docker build -t onnx_gpu .

GPUありとなしで各モデルを実行します。

for MODEL in yolox_nano yolox_tiny yolox_s yolox_m yolox_l yolox_x
do
  docker run --rm -it --gpus all onnx_gpu --model $MODEL
done

for MODEL in yolox_nano yolox_tiny yolox_s yolox_m yolox_l yolox_x
do
  docker run --rm -it onnx_gpu --model $MODEL
done

以下のように出力されました。

// GPUあり
yolox_nano
['CUDAExecutionProvider', 'CPUExecutionProvider']
650it [00:09, 65.09it/s]
9.990914344787598
yolox_tiny
['CUDAExecutionProvider', 'CPUExecutionProvider']
650it [00:09, 71.71it/s]
9.069817543029785
yolox_s
['CUDAExecutionProvider', 'CPUExecutionProvider']
650it [00:11, 54.32it/s]
11.97160267829895
yolox_m
['CUDAExecutionProvider', 'CPUExecutionProvider']
650it [00:16, 40.53it/s]
16.042967319488525
yolox_l
['CUDAExecutionProvider', 'CPUExecutionProvider']
650it [00:21, 30.04it/s]
21.642226457595825
yolox_x
['CUDAExecutionProvider', 'CPUExecutionProvider']
650it [00:33, 19.43it/s]
33.46185088157654

// GPUなし
yolox_nano
['CPUExecutionProvider']
650it [00:10, 63.21it/s]
10.288543939590454
yolox_tiny
['CPUExecutionProvider']
650it [00:17, 37.57it/s]
17.307442903518677
yolox_s
['CPUExecutionProvider']
650it [00:57, 11.37it/s]
57.18255543708801
yolox_m
['CPUExecutionProvider']
650it [02:05,  5.17it/s]
125.62260937690735
yolox_l
['CPUExecutionProvider']
650it [03:59,  2.72it/s]
239.15086388587952
yolox_x
['CPUExecutionProvider']
650it [06:51,  1.58it/s]
411.7909436225891

結果

結果をまとめると以下の通りです。理論値は公式リポジトリからです。

Model Parameters GFLOPs Test Size mAP GPU[s] GPU[fps] CPU[s] CPU[fps]
YOLOX-Nano 0.91M 1.08 416×416 25.8 9.99 65.06 10.29 63.18
YOLOX-Tiny 5.06M 6.45 416×416 32.8 9.07 71.67 17.31 37.56
YOLOX-S 9.0M 26.8 640×640 40.5 11.97 54.30 57.18 11.37
YOLOX-M 25.3M 73.8 640×640 47.2 16.04 40.52 125.62 5.17
YOLOX-L 54.2M 155.6 640×640 50.1 21.64 30.03 239.15 2.72
YOLOX-X 99.1M 281.9 640×640 51.5 33.46 19.43 411.79 1.58

見てわかる通り、GPUありにすると明らかに速くなっています。NanoやTinyは動画の読み込みがボトルネックになって、速度がサチっているので、参考程度です。

CPUの場合は、Flopsと反比例した速度になっていますが、GPUの場合は線形な関係になっていないのが面白かったです。GPUありのYOLOX-Xは20fpsぐらい出ているのが驚き。

このようにすればDocker上でも、GPUありのONNXを使えることがわかりました。ONNXなら環境を統一できますし、いろいろ便利そうですね。

コード

main.py

import argparse
import onnxruntime as ort
import cv2
import time
import numpy as np
from tqdm import tqdm
from yolox_utils import preproc, multiclass_nms, demo_postprocess

def load_session(args):
    print(args.model)

    try:
        provider = ['CUDAExecutionProvider','CPUExecutionProvider']
        session = ort.InferenceSession(f'yolox/{args.model}.onnx', providers=provider)
    except:
        provider = ['CPUExecutionProvider']
        session = ort.InferenceSession(f'yolox/{args.model}.onnx', providers=provider)
    print(session.get_providers())

    if args.model in ["yolox_tiny", "yolox_nano"]:
        input_size = (416, 416)
    else:
        input_size = (640, 640)
    return session, input_size

def main(args):
    session, input_size = load_session(args)
    cap = cv2.VideoCapture("yolox/target_video.mp4")

    start_time = time.time()

    with tqdm() as pbar:
        while cap.isOpened():
            success, frame = cap.read()
            if not success:
                break
            img, ratio = preproc(frame, input_size)

            ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
            output = session.run(None, ort_inputs)
            predictions = demo_postprocess(output[0], input_size, p6=False)[0]

            boxes = predictions[:, :4]
            scores = predictions[:, 4:5] * predictions[:, 5:]

            boxes_xyxy = np.ones_like(boxes)
            boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
            boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
            boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
            boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
            boxes_xyxy /= ratio
            dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)

            pbar.update(1)

    elapsed = time.time() - start_time
    print(elapsed)

    cap.release()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="")
    args = parser.parse_args()

    main(args)

yolox_utils.py

import numpy as np
import cv2

def nms(boxes, scores, nms_thr):
    """Single class NMS implemented in Numpy."""
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= nms_thr)[0]
        order = order[inds + 1]

    return keep

def multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr):
    """Multiclass NMS implemented in Numpy. Class-aware version."""
    final_dets = []
    num_classes = scores.shape[1]
    for cls_ind in range(num_classes):
        cls_scores = scores[:, cls_ind]
        valid_score_mask = cls_scores > score_thr
        if valid_score_mask.sum() == 0:
            continue
        else:
            valid_scores = cls_scores[valid_score_mask]
            valid_boxes = boxes[valid_score_mask]
            keep = nms(valid_boxes, valid_scores, nms_thr)
            if len(keep) > 0:
                cls_inds = np.ones((len(keep), 1)) * cls_ind
                dets = np.concatenate(
                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
                )
                final_dets.append(dets)
    if len(final_dets) == 0:
        return None
    return np.concatenate(final_dets, 0)

def multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr):
    """Multiclass NMS implemented in Numpy. Class-agnostic version."""
    cls_inds = scores.argmax(1)
    cls_scores = scores[np.arange(len(cls_inds)), cls_inds]

    valid_score_mask = cls_scores > score_thr
    if valid_score_mask.sum() == 0:
        return None
    valid_scores = cls_scores[valid_score_mask]
    valid_boxes = boxes[valid_score_mask]
    valid_cls_inds = cls_inds[valid_score_mask]
    keep = nms(valid_boxes, valid_scores, nms_thr)
    if keep:
        dets = np.concatenate(
            [valid_boxes[keep], valid_scores[keep, None], valid_cls_inds[keep, None]], 1
        )
    return dets

def multiclass_nms(boxes, scores, nms_thr, score_thr, class_agnostic=True):
    """Multiclass NMS implemented in Numpy"""
    if class_agnostic:
        nms_method = multiclass_nms_class_agnostic
    else:
        nms_method = multiclass_nms_class_aware
    return nms_method(boxes, scores, nms_thr, score_thr)

def demo_postprocess(outputs, img_size, p6=False):

    grids = []
    expanded_strides = []

    if not p6:
        strides = [8, 16, 32]
    else:
        strides = [8, 16, 32, 64]

    hsizes = [img_size[0] // stride for stride in strides]
    wsizes = [img_size[1] // stride for stride in strides]

    for hsize, wsize, stride in zip(hsizes, wsizes, strides):
        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
        grids.append(grid)
        shape = grid.shape[:2]
        expanded_strides.append(np.full((*shape, 1), stride))

    grids = np.concatenate(grids, 1)
    expanded_strides = np.concatenate(expanded_strides, 1)
    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides

    return outputs

def preproc(img, input_size, swap=(2, 0, 1)):
    if len(img.shape) == 3:
        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
    else:
        padded_img = np.ones(input_size, dtype=np.uint8) * 114

    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
    resized_img = cv2.resize(
        img,
        (int(img.shape[1] * r), int(img.shape[0] * r)),
        interpolation=cv2.INTER_LINEAR,
    ).astype(np.uint8)
    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img

    padded_img = padded_img.transpose(swap)
    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
    return padded_img, r

Dockerfile

FROM nvidia/cuda:11.6.0-cudnn8-runtime-ubuntu20.04

RUN apt-get update
ENV TZ=Asia/Tokyo
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
RUN apt-get install -yq --no-install-recommends python3-pip \
        python3-dev \
        wget \
        git  \
        libopencv-dev \
        tzdata && apt-get upgrade -y && apt-get clean

RUN ln -s /usr/bin/python3 /usr/bin/python

RUN pip install -U pip &&\
  pip install --no-cache-dir onnxruntime-gpu opencv-python tqdm

COPY yolox yolox
COPY src .

ENTRYPOINT ["python", "main.py"]
CMD ["--model", "default"]
このエントリーをはてなブックマークに追加
Pocket
LINEで送る
Delicious にシェア

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

技術書コーナー

北海道の駅巡りコーナー

Tags:, , ,

Add a Comment

メールアドレスが公開されることはありません。 ※ が付いている欄は必須項目です