YOLO AlexeyAB darknet 設定 GPU 編號

最近我同事問我一個問題如何設定 GPU 編號，因為我們平常在使用 GPU 的運算上可能有很多系統需要用到 GPU，會想要把一張顯示卡拆給多個程式使用那要怎麼達成這件事情呢？

查詢 GPU 編號

在開始之前先簡單說明一下，如果想要知道 GPU 編號可以直接使用 nvidia-smi 指令來查詢 GPU 編號，輸入指令後會看到如下的狀態：

$ nvidia-smi
Sun Sep 12 17:41:38 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.73.01    Driver Version: 460.73.01    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  Tesla K80           On   | 00000000:00:1E.0 Off |                    0 |
| N/A   33C    P8    31W / 149W |      0MiB / 11441MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

使用 os.environ 來設定

通常最簡單的方法我們其實可以直接使用 os.environ 來設定 GPU ID，如果你的程式使整隻都給其中一個 GPU 其實用這個方法就可以了，不過如果你的程式希望不同的 function 可以使用不同的 GPU 那這個方法就行不通了。

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"

使用 darknet 提供的 set_gpu

而在 darknet 上面其實有提供一個叫做 set_gpu 的 function 可以讓我們去做指定 GPU 的動作，如果使用此方法就可以分別在不同的檔案或是不同的線程上面使用不一樣的 GPU 了。

我們可以在 darknet.py 的 141 – 143 看到它。

from darknet import set_gpu

set_gpu(1)

參考資料

https://github.com/AlexeyAB/darknet