PyTorch如何切换CPU和GPU的使用详解

PyTorch是一种常用的深度学习框架，它支持在CPU和GPU上运行。在本文中，我们将介绍如何在PyTorch中切换CPU和GPU的使用，并提供两个示例说明。

示例1：在CPU上运行PyTorch模型

以下是一个在CPU上运行PyTorch模型的示例代码：

import torch

# Define model
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Create input tensor
x = torch.randn(1, 10)

# Create model and run on CPU
model = Net()
output = model(x)
print(output)

在这个示例中，我们首先定义了一个简单的神经网络模型。然后，我们创建了一个输入张量，并将模型和输入张量都放在CPU上运行。最后，我们打印了模型的输出。

示例2：在GPU上运行PyTorch模型

以下是一个在GPU上运行PyTorch模型的示例代码：

import torch

# Define model
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Create input tensor
x = torch.randn(1, 10)

# Create model and run on GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
x = x.to(device)
output = model(x)
print(output)

在这个示例中，我们首先定义了一个简单的神经网络模型。然后，我们创建了一个输入张量，并将模型和输入张量都放在GPU上运行。最后，我们打印了模型的输出。

总结

在本文中，我们介绍了如何在PyTorch中切换CPU和GPU的使用，并提供了两个示例说明。这些技术对于在深度学习中处理大规模数据集非常有用。