ONNX RuntimeでGPUメモリー削減につながった意外なこと

またまた ONNX Runtime ネタです。

PyTorchからONNXへの変換を調べるときにまず見るサイトのひとつが、PyTorch Tutorials にある (OPTIONAL) EXPORTING A MODEL FROM PYTORCH TO ONNX AND RUNNING IT USING ONNX RUNTIME かと思います。

無事 model の変換を終えて ONNX Runtime で動作させるときにはきっと、このサンプルコードを参考にするでしょう。

import onnxruntime

ort_session = onnxruntime.InferenceSession("super_resolution.onnx")

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)

# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)

print("Exported model has been tested with ONNXRuntime, and the result looks good!")

ところで、

こうやってモデルを別プラットフォームに変換して運用するとなると、リソース消費の違いや処理速度の違いがとても気になります。

というわけで、今回はGPUメモリーに着目し、以下の3パターンでGPUメモリーの消費量を調べてみたところ、驚きの事実があったのでご紹介したいと思います。

1. PyTorchの場合
2. ONNX Runtimeで上記サンプルコードをコピペした場合
3. ONNX Runtimeで無駄を省きコードを少しきれいにした場合

尚、今回は結果のみのご紹介とさせて頂きます。なぜそうなるのか、さまざまなパターンでの検証などはしておりませんので、あるモデルでの偶然の症状である可能性も十分あることをご承知おきください。

PyTorchの場合

PyTorchでのコードは以下のような感じで、ある時系列データから2種類の時系列データを作って、それをmodelに喰わせてInferenceを行うものです。（いつものMulti-timescale LSTMです）ここではmodelの構築部分は省略していますが、ポイントはmodelがgpuにあることですね。

seq = np.array(<my sequence data>)
x1 = torch.from_numpy(seq[0:t1,:]).unsqueeze(0).float().to('cuda:0')
x2 = torch.from_numpy(seq[t2:t1,:]).unsqueeze(0).float().to('cuda:0')
y_pred = model(x1, x2).to('cpu').squeeze().detach().numpy().copy()

この場合のmodelをloadしたときとmodelを稼働したときのnvidia-smiの結果は以下の通りです。結構複雑なプログラムで検証してしまったので表示も複雑になっていますが、見るべきところは PID 1915465 の GPU Memory Usage です。動作中かどうかは Volatile GPU-Util でわかると思います。

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:01:00.0 Off |                  N/A |
| 55%   50C    P2   104W / 350W |  13956MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+                                                                  
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1268      G   /usr/lib/xorg/Xorg                167MiB |
|    0   N/A  N/A      1526      G   /usr/bin/gnome-shell               11MiB |
|    0   N/A  N/A   1915419      C   ...t_server2/venv/bin/python     8585MiB |
|    0   N/A  N/A   1915463      C   ...t_server2/venv/bin/python     1725MiB |
|    0   N/A  N/A   1915464      C   ...t_server2/venv/bin/python     1737MiB |
|    0   N/A  N/A   1915465      C   ...t_server2/venv/bin/python     1725MiB |
+-----------------------------------------------------------------------------+

$ python3 test.py

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:01:00.0 Off |                  N/A |
| 55%   51C    P2   170W / 350W |  14085MiB / 24268MiB |     39%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+                                  
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1268      G   /usr/lib/xorg/Xorg                167MiB |
|    0   N/A  N/A      1526      G   /usr/bin/gnome-shell               11MiB |
|    0   N/A  N/A   1915419      C   ...t_server2/venv/bin/python     8672MiB |
|    0   N/A  N/A   1915463      C   ...t_server2/venv/bin/python     1733MiB |
|    0   N/A  N/A   1915464      C   ...t_server2/venv/bin/python     1763MiB |
|    0   N/A  N/A   1915465      C   ...t_server2/venv/bin/python     1733MiB |
+-----------------------------------------------------------------------------+

ONNX Runtimeで上記サンプルコードをコピペした場合

PyTorchのコードを元にあまり考えずに書くと以下のような感じになります。(sessionはGPU側で動作するようになっています)
つまり、一旦GPUに持って行った時系列データをまたCPUに戻して、さらにGPUで動作するmodelに転送するという、大変無駄の多いコードになってしまっていますが、ひとまず PyTorch Tutorial を参考に手っ取り早く動作させようとするとこうなりました。

seq = np.array(<my sequence data>)
x1 = torch.from_numpy(seq[0:t1,:]).unsqueeze(0).float().to('cuda:0')
x2 = torch.from_numpy(seq[t2:t1,:]).unsqueeze(0).float().to('cuda:0')

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

inputs = {session.get_inputs()[0].name: to_numpy(x1), session.get_inputs()[1].name: to_numpy(x2)}
y_pred = session.run(None, inputs)[0]

この場合のmodelをloadしたときとmodelを稼働したときのnvidia-smiの結果は以下の通りです。PyTorchでは1.8GB弱だったGPUメモリーの消費が2.6GBまで増加し、残念なことになっていることがわかります。

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:01:00.0 Off |                  N/A |
| 56%   41C    P8    23W / 350W |  13240MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1268      G   /usr/lib/xorg/Xorg                167MiB |
|    0   N/A  N/A      1526      G   /usr/bin/gnome-shell               11MiB |
|    0   N/A  N/A   1930143      C   ...t_server2/venv/bin/python     8585MiB |
|    0   N/A  N/A   1930186      C   ...t_server2/venv/bin/python     1725MiB |
|    0   N/A  N/A   1930187      C   ...t_server2/venv/bin/python     1737MiB |
|    0   N/A  N/A   1930188      C   ...t_server2/venv/bin/python     1009MiB |
+-----------------------------------------------------------------------------+

$ python3 test.py

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:01:00.0 Off |                  N/A |
| 55%   48C    P2   169W / 350W |  14953MiB / 24268MiB |     44%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+                                             
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1268      G   /usr/lib/xorg/Xorg                167MiB |
|    0   N/A  N/A      1526      G   /usr/bin/gnome-shell               11MiB |
|    0   N/A  N/A   1930143      C   ...t_server2/venv/bin/python     8672MiB |
|    0   N/A  N/A   1930186      C   ...t_server2/venv/bin/python     1733MiB |
|    0   N/A  N/A   1930187      C   ...t_server2/venv/bin/python     1763MiB |
|    0   N/A  N/A   1930188      C   ...t_server2/venv/bin/python     2601MiB |
+-----------------------------------------------------------------------------+

ONNX Runtimeで無駄を省きコードを少しきれいにした場合

というわけで、無駄なデータ転送を省いてきれいにしたコードが以下です。

seq = np.array(<my sequence data>)
x1 = (seq[0:t1,:])[np.newaxis].astype(np.float32)
x2 = (seq[t2:t1,:])[np.newaxis].astype(np.float32)
inputs = {session.get_inputs()[0].name: x1, session.get_inputs()[1].name: x2}
y_pred = session.run(None, inputs)[0]

随分すっきりしましたね。この場合のmodelをloadしたときとmodelを稼働したときのnvidia-smiの結果は以下の通りです。たぶん減るだろうとは思っていましたが、1.4GB弱にまで消費量は削減されました。

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:01:00.0 Off |                  N/A |
| 55%   39C    P8    18W / 350W |  13240MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1268      G   /usr/lib/xorg/Xorg                167MiB |
|    0   N/A  N/A      1526      G   /usr/bin/gnome-shell               11MiB |
|    0   N/A  N/A   1904264      C   ...t_server2/venv/bin/python     8585MiB |
|    0   N/A  N/A   1904307      C   ...t_server2/venv/bin/python     1725MiB |
|    0   N/A  N/A   1904308      C   ...t_server2/venv/bin/python     1737MiB |
|    0   N/A  N/A   1904309      C   ...t_server2/venv/bin/python     1009MiB |
+-----------------------------------------------------------------------------+

$ python3 test.py

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:01:00.0 Off |                  N/A |
| 55%   51C    P2   169W / 350W |  13719MiB / 24268MiB |     31%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1268      G   /usr/lib/xorg/Xorg                167MiB |
|    0   N/A  N/A      1526      G   /usr/bin/gnome-shell               11MiB |
|    0   N/A  N/A   1904264      C   ...t_server2/venv/bin/python     8672MiB |
|    0   N/A  N/A   1904307      C   ...t_server2/venv/bin/python     1733MiB |
|    0   N/A  N/A   1904308      C   ...t_server2/venv/bin/python     1763MiB |
|    0   N/A  N/A   1904309      C   ...t_server2/venv/bin/python     1367MiB |
+-----------------------------------------------------------------------------+

でもなぜ ?

GPUメモリーの消費量が減ったのはもちろん喜ばしいことですが、でもなぜここまで削減できたのかというのは今のところ謎です。ただ、この程度のちょっとしたことで結構メモリー消費が変わるというのは、正直驚きですね。

前回の記事でも触れた CUDA Execution Provider にありましたが、providers で gpu_mem_limit を指定できるようなので、ここにいろいろな値を入れてみて試してみるというのはありますね。

とにもかくにも PyTorch よりもGPUメモリーの消費を減らせてよかったです。またCPUで動作させても結構早くなりました。重宝させてもらってます、ONNX Runtime。

関連記事