Tensor 基本操作5 device 管理，使用 GPU 设备 | PyTorch 深度学习实战

Chatopera 研发团队2025-01-27 13:27

前一篇文章，Tensor 基本操作4 理解 indexing，加减乘除和 broadcasting 运算 | PyTorch 深度学习实战

本系列文章 GitHub Repo: https://github.com/hailiang-wang/pytorch-get-started

Tensor 基本使用

检查设备
[创建 tensor 时声明设备](#创建 tensor 时声明设备)
更改默认设备
[创建 tensor 后移动 tensor.to](#创建 tensor 后移动 tensor.to)
注意事项
- [1. 运算的发生位置](#1. 运算的发生位置)
- [2. 当两个 tensor 进行运算时，需要在同一个设备上](#2. 当两个 tensor 进行运算时，需要在同一个设备上)

检查设备

默认为 CPU，根据是否有 GPU 设定 device。

device = (torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
检查默认设备

device = torch.get_default_device()

创建 tensor 时声明设备

复制代码

    # 声明为 CPU 设备
    device_cpu = torch.device('cpu')
    points_cpu = torch.tensor([[4.0, 1.0], [5.0, 3.0], [2.0, 1.0]], device=device_cpu)

更改默认设备

将默认设备，设置为 GPU。

复制代码

device_gpu = torch.device('cuda')
torch.set_default_device(device_gpu) 
points_default = torch.tensor([[4.0, 1.0], [5.0, 3.0], [2.0, 1.0]]) # 此时，points_default  被定义到 GPU 上

创建 tensor 后移动 tensor.to

将一个 tensor 移动到指定设备。

复制代码

    device_gpu = torch.device('cuda')
    points2 = points.to(device_gpu)  # 将 Tensor 复制到 GPU
    print(points2)

注意事项

1. 运算的发生位置

复制代码

    device_gpu = torch.device('cuda')
    points2 = points.to(device_gpu)  # 将 Tensor 复制到 GPU
    points3 = points2 * 2   # points3 还是在 GPU 上
    points4 = points2 + 2  # points4 还是在 GPU 上

但是，打印 points3 或 points4 时，将会复制该值到 CPU 上输出。

2. 当两个 tensor 进行运算时，需要在同一个设备上

复制代码

  File "C:\devel\Python\Python311\Lib\site-packages\torch\utils\_device.py", line 79, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!