PyTorch中的CUDA的操作方法

2022-08-26 14:02:58

前言

CUDA(Compute Unified Device Architecture)是NVIDIA推出的異構計算平臺，PyTorch中有專門的模組torch.cuda來設定和執行CUDA相關操作。本地安裝環境為Windows10，Python3.7.8和CUDA 11.6，安裝PyTorch最新穩定版本1.12.1如下：

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

一.常見CPU和GPU操作命令

1.檢視PyTorch版本

print(torch.__version__)
1.12.1+cu116

2.檢視GPU裝置是否可用

print(torch.cuda.is_available())
True

3.PyTorch預設使用裝置是CPU

print("default device: {}".format(torch.Tensor([4,5,6]).device))
default device: cpu

4.檢視所有可用的cpu裝置的數量

print("available cpu devices: {}".format(torch.cuda.os.cpu_count()))
available cpu devices: 20

這裡CPU裝置數量指的是邏輯處理器的數量。

5.檢視所有可用的gpu裝置的數量

print("available gpu devices: {}".format(torch.cuda.device_count()))
available gpu devices: 1

6.獲取gpu裝置的名稱

print("gpu device name: {}".format(torch.cuda.get_device_name(torch.device("cuda:0"))))
gpu device name: NVIDIA GeForce GTX 1080 Ti

7.通過device="cpu:0"指定cpu:0裝置

device = torch.Tensor([1,2,3], device="cpu:0").device
print("device type: {}".format(device))
device type: cpu

8.通過torch.device指定cpu:0裝置

cpu1 = torch.device("cpu:0")
print("cpu device: {}:{}".format(cpu1.type, cpu1.index))
cpu device: cpu:0

9.使用索引的方式，預設使用CUDA裝置

gpu = torch.device(0)
print("gpu device: {}:{}".format(gpu.type, gpu.index))
gpu device: cuda:0

10.通過torch.device("cuda:0)指定cuda:0裝置

gpu = torch.device("cuda:0")
print("gpu device: {}:{}".format(gpu.type, gpu.index))
gpu device: cuda:0

二.CPU和GPU裝置上的Tensor

預設情況下建立Tensor是在CPU裝置上的，但是可以通過copy_、to、cuda等方法將CPU裝置中的Tensor轉移到GPU裝置上。當然也是可以直接在GPU裝置上建立Tensor的。torch.tensor和torch.Tensor的區別是，torch.tensor可以通過device指定gpu裝置，而torch.Tensor只能在cpu上建立，否則報錯。

1.Tensor從CPU拷貝到GPU上

# 預設建立的tensor是在cpu上建立的
cpu_tensor = torch.Tensor([[1,4,7],[3,6,9],[2,5,8]])
print(cpu_tensor.device)

# 通過to方法將cpu_tensor拷貝到gpu上
gpu_tensor1 = cpu_tensor.to(torch.device("cuda:0"))
print(gpu_tensor1.device)

# 通過cuda方法將cpu_tensor拷貝到gpu上
gpu_tensor2 = cpu_tensor.cuda(torch.device("cuda:0"))
print(gpu_tensor2.device)

# 將gpu_tensor2拷貝到cpu上
gpu_tensor3 = cpu_tensor.copy_(gpu_tensor2)
print(gpu_tensor3.device)
print(gpu_tensor3)

輸出結果如下：

cpu
cuda:0
cuda:0
cpu
tensor([[1., 4., 7.],
[3., 6., 9.],
[2., 5., 8.]])

主要說明下這個copy_()方法，實現如下：

def copy_(self, src, non_blocking=False):
    ......
    return _te.Tensor(*(), **{})

就是從src中拷貝元素到self的tensor中，然後返回self。以gpu_tensor3 = cpu_tensor.copy_(gpu_tensor2)為例，就是把gpu中的gpu_tensor2拷貝到cpu中的cpu_tensor中。

2.直接在GPU上建立Tensor

gpu_tensor1 = torch.tensor([[2,5,8],[1,4,7],[3,6,9]], device=torch.device("cuda:0"))
print(gpu_tensor1.device)

# 在gpu裝置上建立亂數tensor
print(torch.rand((3,4), device=torch.device("cuda:0")))

# 在gpu裝置上建立0值tensor
print(torch.zeros((2,5), device=torch.device("cuda:0")))

輸出結果，如下：

cuda:0
tensor([[0.7061, 0.2161, 0.8219, 0.3354],
[0.1697, 0.1730, 0.1400, 0.2825],
[0.1771, 0.0473, 0.8411, 0.2318]], device='cuda:0')
tensor([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], device='cuda:0')

3.CUDA Streams

Steam是CUDA命令線性執行的抽象形式，分配給裝置的CUDA命令按照入隊序列的順序執行。每個裝置都有一個預設的Steam，也可以通過torch.cuda.Stream()建立新的Stream。如果不同Stream中的命令互動執行，那麼就不能保證命令絕對按順序執行。下面的這個例子不同的Stream就可能會產生錯誤。

cuda = torch.device("cuda")
# 建立預設的stream，A就是使用的預設stream
s = torch.cuda.Stream()
A = torch.randn((1,10), device=cuda)
for i in range(100):
    # 在新的stream上對預設的stream上建立的tensor進行求和
    with torch.cuda.stream(s):
        # 存在的問題是：torch.sum()可能會在torch.randn()之前執行
        B = torch.sum(A)
        print(B)

這個例子存在的問題是torch.sum()可能會在torch.randn()之前就執行。為了保證Stream中的命令絕對按順序執行，接下來使用Synchronize同步方法解決上面例子的問題：

cuda = torch.device("cuda")
s = torch.cuda.Stream()
A = torch.randn((1,10), device=cuda)
default_stream = torch.cuda.current_stream()
print("Default Stream: {}".format(default_stream))
# 等待建立A的stream執行完畢
torch.cuda.Stream.synchronize(default_stream)
for i in range(100):
    # 在新的stream上對預設的stream上建立的tensor進行求和
    with torch.cuda.stream(s):
        print("current stream: {}".format(torch.cuda.current_stream()))
        B = torch.sum(A)
        print(B)

解決問題的思路就是通過torch.cuda.Stream.synchronize(default_stream)等待建立A的stream執行完畢，然後再執行新的Stream中的指令。
除此之外，使用memory_cached方法獲取快取記憶體的大小，使用max_memory_cached方法獲取最大快取記憶體的大小，使用max_memory_allocated方法獲取最大分配記憶體的大小。可以使用empty_cache方法釋放無用的快取記憶體。

三.固定緩衝區

快取就是當計算機記憶體不足的時候，就會把記憶體中的資料儲存到硬碟上。固定緩衝區就是說常駐記憶體，不能把這部分資料快取到硬碟上。可以直接使用pin_memory方法或在Tensor上直接呼叫pin_memory方法將Tensor複製到固定緩衝區。為什麼要做固定緩衝區呢？目的只有一個，就是把CPU上的固定緩衝區拷貝到GPU上時速度快。Tensor上的is_pinned方法可以檢視該Tensor是否載入到固定緩衝區中。

from torch.utils.data._utils.pin_memory import pin_memory
x = torch.Tensor([[1,2,4], [5, 7, 9], [3, 7, 10]])
# 通過pin_memory()方法將x複製到固定緩衝區
y = pin_memory(x)
# 在tensor上直接呼叫pin_memory()方法將tensor複製到固定緩衝區
z = x.pin_memory()
# id()方法返回tensor的記憶體地址，pin_memory()返回tensor物件的拷貝，因此記憶體地址是不同的
print("id: {}".format(id(x)))
print("id: {}".format(id(y)))
print("id: {}".format(id(z)))
# 當tensor放入固定緩衝區後，就可以非同步將資料複製到gpu裝置上了
a = z.cuda(non_blocking=True)
print(a)
print("is_pinned: {}/{}".format(x.is_pinned(), z.is_pinned()))

輸出結果如下所示：

id: 1605289350472
id: 1605969660408
id: 1605969660248
tensor([[ 1., 2., 4.],
[ 5., 7., 9.],
[ 3., 7., 10.]], device='cuda:0')
is_pinned: False/True

說明：通過id()檢視物件的記憶體地址。

四.自動裝置感知

1.適配CPU和GPU裝置

自動裝置感知本質上就是有GPU時就使用GPU，沒有GPU時就使用CPU，即一套程式碼適配CPU和GPU裝置。GPU是否存在是通過torch.cuda.is_available()判斷的。

常見的寫法如下：

device = torch.device("cpu")
if torch.cuda.is_available():
    device = torch.device("cuda")
a = torch.tensor([1,2,3], device=device)
print(a)

輸出結果如下所示：

tensor([1, 2, 3], device='cuda:0')

2.模型遷移到GPU裝置

在Module物件上呼叫to()方法可以把模型也遷移到GPU裝置上，如下所示：

class LinearRegression(torch.nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = torch.nn.Linear(1, 1)
    def forward(self, x):
        return self.linear(x)
regression = LinearRegression().to(device=device)
for param in regression.parameters():
    print(param)

從上述輸出引數中可以看到param都是device='cuda:0’上的tensor，所以可以說模型通過to()遷移到GPU裝置上了。

到此這篇關於PyTorch中的CUDA的操作方法的文章就介紹到這了,更多相關PyTorch CUDA操作內容請搜尋it145.com以前的文章或繼續瀏覽下面的相關文章希望大家以後多多支援it145.com！