Pytorch 如何默认在GPU上运行PyTorch

在本文中，我们将介绍如何通过默认设置在GPU上运行PyTorch。PyTorch是一个基于Torch的Python机器学习库，提供了强大的GPU加速功能，可以大大提高训练和推理的速度。然而，默认情况下，PyTorch会在CPU上运行，所以我们需要进行一些设置才能将其切换到GPU上。

阅读更多：Pytorch 教程

检测GPU是否可用

首先，我们需要检查系统是否安装了合适的GPU驱动和CUDA工具包，并且PyTorch是否已经正确安装。我们可以使用下面的代码来检测GPU是否可用：

import torch

if torch.cuda.is_available():
    print("GPU is available!")
else:
    print("GPU is not available.")

如果输出结果为”GPU is available!”，说明系统已经正确安装了GPU驱动和CUDA工具包，并且PyTorch可以在GPU上运行。如果输出结果为”GPU is not available.”，则需要检查系统是否符合GPU运行的要求，并进行相应的配置。

将PyTorch切换到GPU上

接下来，我们需要将PyTorch切换到GPU上。PyTorch提供了一个torch.device对象来管理设备的选择。我们可以使用以下代码将PyTorch切换到GPU上：

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型和数据移动到设备上
model.to(device)
data = data.to(device)

在上面的代码中，我们首先通过torch.cuda.is_available()检查GPU是否可用，然后根据结果选择设备：”cuda”表示使用GPU，”cpu”表示使用CPU。然后，我们可以使用.to(device)方法将模型、数据等对象移动到选择的设备上。

默认设置为GPU

为了方便起见，我们希望将PyTorch默认设置为GPU。通过设置环境变量CUDA_VISIBLE_DEVICES，我们可以指定使用哪些GPU设备。

export CUDA_VISIBLE_DEVICES=0,1

上述命令会将可见的GPU设备设置为0和1。这样，在代码中不需要显式地指定设备，PyTorch将自动使用指定的GPU设备。

示例说明

下面是一个完整的示例代码，演示了如何将PyTorch默认设置为GPU并进行训练：

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载训练集和测试集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.ReLU()(x)
        x = self.conv2(x)
        x = nn.ReLU()(x)
        x = nn.MaxPool2d(2)(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.ReLU()(x)
        x = self.fc2(x)
        return x

# 创建模型实例并将其移动到设备上
net = Net()
net.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练网络
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 100 == 99:
            print(f"[{epoch + 1}, {i + 1}] loss: {running_loss / 100:.3f}")
            running_loss = 0.0

print("Finished training")

# 测试网络
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%")

在上面的示例代码中，我们首先使用torchvision加载MNIST数据集，并进行了一些基本的数据预处理。然后，我们定义了一个简单的卷积神经网络，并将其移动到设备上。接下来，我们定义了损失函数和优化器，并进行了训练和测试。

总结

通过本文的介绍，我们学习了如何通过默认设置在GPU上运行PyTorch。我们首先检测系统是否支持GPU，并通过torch.device对象将PyTorch切换到GPU上。然后，我们了解了如何设置环境变量CUDA_VISIBLE_DEVICES来将PyTorch默认设置为GPU。最后，我们通过一个示例演示了如何在GPU上训练一个简单的卷积神经网络。

在实际中，使用GPU可以大大加快深度学习的训练和推理速度，尤其是当处理大规模数据集或复杂模型时。因此，了解如何将PyTorch默认设置为GPU可以帮助我们充分利用GPU的性能优势。希望本文对您有所帮助！