使用 PyTorch 進行 CUDA 程式開發在輝達 nVidia Jetson Nano 環境下

在硬體上，我們有了 Jetson Orin Nano 的多核心 GPU 的裝置之後；接下來可以利用 PyTorch 來作 CUDA python 程式開發。由於PyTorch 的語法跟常用的 numpy 程式庫語法很接近，我自己認為會比較適合非資訊專業的人員作為 CUDA 多核心 GPU 程式開發及研究的工具。

• 確認 CUDA 版本以及是否可以使用

我們可以檢查一下，PyTorch 跟 CUDA 是否可以順利執行；如果我們按照 nVidia 提供的 PyTorch 檔案套件來安裝的話，應該可以看到 PyTorch 的版本是「xx.xx.xx+nvxx.xx」。例如：

import torch
# check torch version
print(‘pytorch version’,torch.__version__)
# check available device
if torch.cuda.is_available():
	print(‘cuda device:’,torch.cuda.get_device_name())

• PyTorch 的張量 (tensor) 資料型態

在 PyTorch 編寫 CUDA 的程式時，主要的計算需要把數字轉換成「張量」(Tensor) 的形態；我們可以用以下的程式來大致了解，在一般定義上，「張量」和「矩陣」以及「純量」運算的差異。

import torch
import numpy as np
#------------
# scale
#------------
Scale=1
#------------
# matrix
#------------
Matrix=np.array([[1,2],[3,4]])
#------------
# tensor
#------------
Tensor=torch.tensor([[[1,2],[3,4]],[[5,6],[7,8]]])
#------------
print('Scale:',Scale)
print('Matrix:',Matrix)
print('Tensor:',Tensor)
print('Matrix: dimension is',Matrix.ndim,'size is',Matrix.shape)
print('Tensor: dimension is',Tensor.ndim,'size is',Tensor.shape)

不過，在 PyTorch 使用上，「torch.tensor()」這個函數其實是可以同時來使用「張量」和「矩陣」以及「純量」；例如：

import torch
#------------
# scale
#------------
Scale=torch.tensor([1])
#------------
# matrix
#------------
Matrix=torch.tensor([[1,2],[3,4]])
#------------
# tensor
#------------
Tensor=torch.tensor([[[1,2],[3,4]],[[5,6],[7,8]]])
#------------
print('Scale:',Scale)
print('Matrix:',Matrix)
print('Tensor:',Tensor)
print('Scale: dimension is',Scale.ndim,',size is',Scale.shape)
print('Matrix: dimension is',Matrix.ndim,',size is',Matrix.shape)
print('Tensor: dimension is',Tensor.ndim,',size is',Tensor.shape)

• 撰寫在 GPU 上運算的程式

在利用 PyTorch 來使用 CUDA 計算時，要將數值轉換成「tensor」的型態，然後指定到要計算的 GPU 上面，再利用 PyTorch 的提供的函數來作運算。下例中，我們使用 PyTorch/CUDA 在多 GPU 核心上，作矩陣相乘的運算；要注意的是，CUDA 的運算是採用非同步運算，如果在 CUDA 運算完之後，要將結果傳回 CPU 端來輸出的話，最好是要使用強迫同步的指令。

import numpy as np
import torch
import time

# 設定亂數矩陣
a=np.random.rand(1000,1000)
b=np.random.rand(1000,1000)
# 指定到 CUDA 作計算
device = torch.device('cuda')
valueA=torch.tensor(a).to(device)
valueB=torch.tensor(b).to(device)
# 開始計算時間
beginTime=time.time()
for i in range(100):
    torch.matmul(valueA,valueB)
# 等所有 GPU 都計算完畢
torch.cuda.synchronize()
# 印出計算所花的時間
print('CUDA execution time:',time.time()-beginTime)

在這個程式在 Jupyter 上撰寫並運算時，我們可以同時開啟 Jetson Orin Nano 的「jtop」程式，確認 Jetson Orin Nano 的 GPU 是否在正在進行運算；如下圖。

這樣子，我們就可以透過 PyTorch 及 CUDA 在 Jetson Orin Nano 的裝置上，進行多 GPU 核心的程式開發了。