GPU加速PyTorch训练：CUDA与MPS的使用指南

GPU加速PyTorch训练：CUDA与MPS的使用指南

引用

CSDN

1.

https://blog.csdn.net/yanzhenjie1003/article/details/136665876

在深度学习领域，神经网络模型的训练需要庞大的算力。GPU凭借其强悍的并行计算能力和内存带宽，已经成为业界首选的解决方案。本文将详细介绍如何使用PyTorch进行GPU加速训练，包括CUDA和MPS的使用方法以及CPU和GPU之间的切换。

1、GPU 的并行计算

深度学习主要包括训练（Training）和推理（Inference）两个环节。训练环节通过投喂大量数据来训练复杂的神经网络模型，而推理环节则利用训练好的模型进行数据推理。训练环节由于涉及海量数据和复杂神经网络结构，对芯片算力性能要求极高。

从硬件角度来看，CPU和GPU的角色有所不同：

• CPU：像一个拥有几十名高级厨师的全能型餐厅，能够处理各种复杂的计算任务，但处理速度相对较慢。
• GPU：像一个拥有千万名初级厨师的单一型餐厅，专门处理特定类型的计算任务，由于并行计算能力强，处理速度非常快。

图一：CPU和GPU的角色对比

GPU凭借其强大的并行计算能力和内存带宽，可以很好地应对深度学习中的训练和推理任务。一块优化良好的GPU可以提供相当于几十块CPU的算力。

2、CUDA 和 MPS

在PyTorch中，torch.cuda和torch.mps是两个用于加速训练的重要模块。在安装PyTorch时，需要根据自己的电脑配置选择正确的版本。

2.1、非 M 芯片验证 CUDA 可用性

对于非M系列芯片的电脑（如OS X、Linux、Windows），安装命令通常包含CUDA：

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch
conda install pytorch torchvision torchaudio cudatoolkit -c pytorch

可以通过以下代码验证GPU状态：

import torch
print(torch.cuda.is_available())

如果GPU可用，将输出True，否则输出False。

2.2、M 芯片验证 MPS 可用性

对于M系列芯片，无法安装CUDA。从PyTorch v1.12版本开始，PyTorch支持了Mac OS的GPU加速，通过苹果的MPS（Metal Performance Shaders）来实现。MPS针对每个Metal GPU系列的独特特性进行了优化。

可以通过以下代码验证MPS是否可用：

import torch
print(torch.backends.mps.is_available())

如果MPS可用，将输出True，否则输出False。

3、CPU 和 GPU 的切换

要使用GPU加速训练，首先需要确认设备上至少有一个GPU可用。然后，需要将数据从CPU内存（RAM）移动到GPU内存（VRAM）。

默认情况下，创建的张量是在CPU上的：

import torch
x = torch.Tensor([1, 2, 3])
print(x, x.device)

输出：

tensor([1., 2., 3.]) cpu

可以使用以下代码创建在GPU上运行的张量：

if torch.cuda.is_available():
    my_device = torch.device('cuda')
elif torch.backends.mps.is_available():
    my_device = torch.device('mps')
else:
    my_device = torch.device('cpu')
print('Device: {}'.format(my_device))
x = torch.rand(2, 2, device=my_device)
print(x)

如果GPU支持CUDA，输出：

Device: cuda
tensor([[0.0024, 0.6778],
        [0.2441, 0.6812]], device='cuda:0')

如果GPU支持MPS，输出：

Device: mps
tensor([[0.5986, 0.4086],
        [0.0624, 0.7131]], device='mps:0')

可以使用.to函数在CPU和GPU之间切换张量：

device_name = (
    "cuda" if torch.cuda.is_available() else "mps"
    if torch.backends.mps.is_available() else "cpu"
)
cpu_device = torch.device("cpu")
gpu_device = torch.device(device_name)
x = torch.Tensor([1, 2, 3])
x = x.to(cpu_device)
print(x, x.device)
x = x.to(gpu_device)
print(x, x.device)

输出：

tensor([1., 2., 3.]) cpu
tensor([1., 2., 3.], device='mps:0') mps:0

通过以上步骤，可以充分利用GPU的并行计算能力，加速PyTorch的训练过程。