DockerでGPU版ONNXを使ってみる
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"]
Shikoan's ML Blogの中の人が運営しているサークル「じゅ~しぃ~すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー