PyTorch中的CUDA的操作方法

Posted on by

前言

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操作內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet！

推薦閱讀：