PyTorch_指定运算设备 （包含安装 GPU 的 PyTorch）

PyTorch默认会将张量创建在 CPU 控制的内存中，即：默认的运算设备为 CPU。我们也可以将张量创建在 GPU 上，能够利用对于矩阵计算的优势加快模型训练。

将张量移动到 GPU 上有两种方法：

1. 使用 cuda 方法
2. 直接在 GPU 上创建张量
3. 使用 to 方法指定设备

安装含有 GPU 的 PyTorch

通过这个可以判断电脑里是否已经安装 CUDA

import torch 

print(torch.cuda.is_available()) # 判断是否有可用的 GPU 设备

如果结果输出是 False，说明设备里没有 GPU 的 PyTorch 的版本。

可以通过 PyTorch官网来安装含有 GPU 的 PyTorch 的版本。

安装 CUDA 12.8 版本的pip命令

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

这样就安装好含有 GPU 的 PyTorch 了。

代码

import torch 

# 使用 cuda 方法
def test01():
    data = torch.tensor([10, 20, 30])
    print("存储设备：", data.device)

    # 将张量移动到 GPU 设备上
    data = data.cuda()
    print("存储设备：", data.device)
    
    # 将张量从 GPU 移动到 CPU 设备上
    data = data.cpu()
    print("存储设备：", data.device)

# 直接将张量创建在指定设备上
def test02():
    data = torch.tensor([10, 20, 30], device='cuda')
    print("存储设备：", data.device)

    # 将张量移动到 CPU 设备上
    data = data.cpu()
    print("存储设备：", data.device)

# 使用 to 方法
def test03():
    data = torch.tensor([10, 20, 30])
    print("存储设备：", data.device)

    # 将张量移动到 GPU 设备上
    data = data.to('cuda')
    print("存储设备：", data.device)
    
    # 将张量从 GPU 移动到 CPU 设备上
    data = data.to('cpu')
    print("存储设备：", data.device)


# 注意点：张量存储在不同设备上的张量不能够直接运算
def test04():
    data1 = torch.tensor([10,20,30])
    data2 = torch.tensor([1,2,3], device='cuda')

    # RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
    # data1 = data1.to('cuda') # 这样可以解决了
    data = data1 + data2 
    print(data)

    # 如果你的电脑上安装 pytorch 不是 gpu 版本的，或者电脑本身没有 gpu 设备环境
    # 否则下面的调用 cuda 函数的代码会报错
    # data1 = data1.cuda() 

if __name__ == "__main__":
    test04()