Pytorch如何指定device(cuda or cpu)例子解析

乔丹搞IT2024-08-26 10:08

代码示例：

在PyTorch中，指定设备（CPU或CUDA）是一个非常重要的步骤，特别是当你在进行深度学习训练时。以下是一些指定设备的详细例子：

检查CUDA是否可用 :

首先，你需要检查你的机器是否支持CUDA，并且PyTorch是否能够使用CUDA。

python 复制代码

import torch
if torch.cuda.is_available():
    print("CUDA is available. Using GPU.")
else:
    print("CUDA is not available. Using CPU.")

设置默认设备 :

你可以设置PyTorch的默认设备，这样所有的张量和模型都会默认使用这个设备。
python 复制代码
```
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
```

创建张量并指定设备 :

当你创建张量时，可以指定它们应该在哪个设备上。

python 复制代码

# 创建一个在CPU上的张量
x_cpu = torch.tensor([1., 2., 3.], device='cpu')

# 创建一个在GPU上的张量
x_gpu = torch.tensor([1., 2., 3.], device=device)

将张量移动到指定设备 :

如果张量已经创建，你可以使用.to()或.cuda()方法将其移动到指定的设备。

python 复制代码

# 将张量移动到GPU
x_gpu = x_cpu.to(device)

# 如果你知道你的设备是GPU，也可以使用.cuda()
if torch.cuda.is_available():
    x_gpu = x_cpu.cuda()

指定模型的设备 :

当你定义模型时，可以将其放置在指定的设备上。
python 复制代码
```
model = MyModel().to(device)
```

在训练循环中使用设备 :

在训练循环中，你需要确保模型的输入数据和目标也在正确的设备上。

python 复制代码

for data, target in dataloader:
    data, target = data.to(device), target.to(device)

    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

保存和加载模型时指定设备 :

当你保存或加载模型时，确保模型在正确的设备上。

python 复制代码

# 保存模型
torch.save(model.state_dict(), "model.pth")

# 加载模型
model = MyModel()
model.load_state_dict(torch.load("model.pth", map_location=device))
model.to(device)

请注意，当你在GPU上训练时，所有的输入数据、目标、模型参数等都应该在GPU上。这样可以确保计算是在GPU上进行的，从而提高训练速度。如果你的机器有多个GPU，你还可以指定使用特定的GPU，例如：

python 复制代码

device = torch.device("cuda:0")  # 使用第一个GPU

以上就是在PyTorch中指定设备的一些基本方法和例子。

喜欢本文，请点赞、收藏和关注！