如何使用 PyTorch 脚本识别 GPU

PyTorch 是一款流行的深度学习框架，可以通过 GPU 快速地训练模型。在使用 PyTorch 时，我们可以利用 CUDA 加速模型的训练和推理，而要想使用 CUDA，首先需要明确当前计算机是否支持 GPU。

本文将介绍如何在 PyTorch 脚本中识别 GPU，并使用 GPU 运行模型。

检测是否安装 CUDA

在开始之前，我们需要确认 CUDA 是否已经安装。我们可以通过以下命令检测 CUDA 是否已经安装：

$ nvcc --version

# 输出类似以下信息：
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2021 NVIDIA Corporation
# Built on Sun_Feb_14_21:12:58_PST_2021
# Cuda compilation tools, release 11.2, V11.2.152
# Build cuda_11.2.r11.2/compiler.29618528_0

如果输出了类似 command not found 的提示，那么说明当前机器上未安装 CUDA。

如果您需要安装 CUDA，请参照官方文档进行安装，安装完成后重新运行上述命令。

检测 GPU 是否可用

我们可以使用以下代码片段检测当前计算机是否支持使用 GPU：

import torch

if torch.cuda.is_available():
    print('GPU is available!')
else:
    print('GPU is NOT available')

如果输出了 GPU is available!，则说明您的计算机可用 GPU 运行 PyTorch。

使用 GPU 运行模型

如果您已经确认当前计算机支持使用 GPU，下面是如何在 PyTorch 中使用 GPU 运行模型。

import torch

# 创建模型
model = ...

# 将模型参数发送到 GPU
model.cuda()

# 定义损失函数
criterion = ...

# 定义优化器
optimizer = ...

# 训练神经网络
for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        # 发送数据到 GPU
        inputs, targets = inputs.cuda(), targets.cuda()

        # 前向传播
        outputs = model(inputs)

        # 计算损失
        loss = criterion(outputs, targets)

        # 反向传播并更新参数
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在上述代码中，我们需要将数据和模型参数都发送到 GPU 中，这样才能够加速训练。在数据发送到 GPU 后，我们就可以像使用 CPU 一样进行训练。

结论

在使用 PyTorch 时，我们可以利用 GPU 对深度学习模型进行训练和推理加速。本文介绍了如何在 PyTorch 脚本中识别 GPU，并使用 GPU 运行模型。

希望这篇文章能够帮助您更好地使用 PyTorch。