深度学习计算,特别是涉及大型张量和复杂模型的计算,需要强大的计算能力。中央处理器(CPU)用途广泛,而图形处理器(GPU)提供大规模并行能力,可以大幅加速构成神经网络核心的矩阵和向量运算。PyTorch 提供简单直接的机制来管理张量的存放位置和计算发生地。了解如何在 CPU 和 GPU 之间移动张量是高效模型训练和推理的必备技能。CPU:默认计算中心默认情况下,当您创建 PyTorch 张量时未指定设备,它会分配在 CPU 上。import torch # 默认在 CPU 上创建张量 cpu_tensor = torch.tensor([1.0, 2.0, 3.0]) print(f"默认张量设备:{cpu_tensor.device}")CPU 完全适用于许多任务,包括预处理步骤、较小规模的计算,或在没有兼容 GPU 时运行模型。然而,对于训练大型深度学习模型,由于 CPU 的顺序处理特性与 GPU 的并行架构相比,仅依赖 CPU 常常会导致训练时间过长,难以承受。GPU:使用并行加速深度学习GPU 包含数百或数千个核心,旨在同时执行大量计算。这种架构非常适合深度学习中常见的运算类型,如大型矩阵乘法和卷积。PyTorch 通过 CUDA(统一计算设备架构)平台使用 NVIDIA GPU。要使用 GPU,您需要:一个兼容 CUDA 的 NVIDIA GPU。已安装适当的 CUDA 工具包。安装了 CUDA 支持的 PyTorch 版本。检查 GPU 可用性并设置设备在尝试使用 GPU 之前,最好检查它是否可用且已为 PyTorch 正确配置。torch.cuda.is_available() 函数在 PyTorch 可以访问支持 CUDA 的 GPU 时返回 True。之后我们可以创建一个 torch.device 对象来表示我们的目标计算设备(CPU 或 GPU)。这使得代码具有适应性,如果 GPU 可用则自动使用 GPU,否则回退到 CPU。import torch # 检查 CUDA 可用性并相应设置设备 if torch.cuda.is_available(): device = torch.device("cuda") # 使用第一个可用的 CUDA 设备 print(f"CUDA (GPU) 可用。使用设备:{device}") # 您也可以指定特定 GPU,例如 torch.device("cuda:0") else: device = torch.device("cpu") print(f"CUDA (GPU) 不可用。使用设备:{device}") # device 现在包含 torch.device('cuda') 或 torch.device('cpu')直接在设备上创建张量您可以在张量创建期间使用 device 参数直接指定目标设备。这通常比在 CPU 上创建然后再移动更高效。# 直接在选定设备上创建张量 try: # 如果 device='cpu',此张量将在 CPU 上;如果 device='cuda',则在 GPU 上 device_tensor = torch.randn(3, 4, device=device) print(f"张量创建于:{device_tensor.device}") except RuntimeError as e: print(f"无法直接在 {device} 上创建张量:{e}") # 处理未找到 GPU 等情况在 CPU 和 GPU 之间移动张量通常,您需要在设备之间传输现有张量。例如,从磁盘加载的数据通常位于 CPU 上,但您的模型可能在 GPU 上以加快计算。移动张量的主要方法是 .to() 方法。.to() 方法接受 torch.device 对象、设备字符串(例如 'cuda'、'cpu'),甚至另一个张量(在这种情况下,张量会移动到与参数张量相同的设备上)作为输入。它在指定设备上返回一个 新 张量。原始张量在其原始设备上保持不变。# 从 CPU 张量开始 cpu_tensor = torch.ones(2, 2) print(f"原始张量:{cpu_tensor.device}") # 将张量移动到选定设备(如果 GPU 可用,则为 GPU;否则为 CPU) # 请记住,'device' 是根据可用性预先设置的 moved_tensor = cpu_tensor.to(device) print(f"移动后的张量:{moved_tensor.device}") # 如果张量在 GPU 上,则显式移回 CPU if moved_tensor.is_cuda: # 检查张量是否在 CUDA 设备上 back_to_cpu = moved_tensor.to("cpu") print(f"张量移回至:{back_to_cpu.device}")PyTorch 还提供便利方法:.cpu() 和 .cuda()。它们分别是 .to('cpu') 和 .to('cuda:0')(或当前默认的 CUDA 设备)的简写。# 使用便利方法(假设 GPU 可用且 'device' 为 'cuda') if device.type == 'cuda': # 将 cpu_tensor 移动到 GPU gpu_tensor_alt = cpu_tensor.cuda() print(f"使用 .cuda():{gpu_tensor_alt.device}") # 将 gpu_tensor_alt 移回 CPU cpu_tensor_alt = gpu_tensor_alt.cpu() print(f"使用 .cpu():{cpu_tensor_alt.device}")设备管理的重要考量设备一致性: 涉及多个张量的操作(例如,加法、矩阵乘法)通常要求所有参与的张量都在 同一 设备上。尝试在 CPU 张量和 GPU 张量之间进行操作将导致 RuntimeError。在执行操作前,请确保您的数据和模型位于同一设备上。# 错误示例(假设 device='cuda') cpu_a = torch.randn(2, 2) gpu_b = torch.randn(2, 2, device=device) try: # 如果 device 是 'cuda',这很可能会失败 c = cpu_a + gpu_b except RuntimeError as e: print(f"在不同设备上执行操作时出错:{e}")数据传输开销: 在 CPU 和 GPU 内存之间移动数据并非瞬间完成。虽然 GPU 计算速度快,但如果管理不当,数据的来回传输可能成为瓶颈。为了最佳性能,请尝试在 GPU 上尽可能多地执行操作,然后在需要时(例如,保存到磁盘或转换为 NumPy)再将最终结果移回 CPU。模型放置: 与张量一样,使用 torch.nn.Module 定义的神经网络模型也需要使用 .to(device) 方法移动到适当的设备。这确保模型的参数(本身也是张量)位于目标设备上进行计算。在讨论模型构建时,会更详细地介绍这一点。掌握使用 device 和 .to() 方法进行张量放置是发挥 GPU 计算能力和编写高效、硬件感知型 PyTorch 代码不可或缺的要素。请记住始终检查设备一致性并注意数据传输成本。