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

代码示例：

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

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

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

2. 设置默认设备:
   你可以设置PyTorch的默认设备，这样所有的张量和模型都会默认使用这个设备。

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

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

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

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

   # 将张量移动到GPU
   x_gpu = x_cpu.to(device)
   
   # 如果你知道你的设备是GPU，也可以使用.cuda()
   if torch.cuda.is_available():
       x_gpu = x_cpu.cuda()
   

5. 指定模型的设备:
   当你定义模型时，可以将其放置在指定的设备上。

   model = MyModel().to(device)
   

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

   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()
   

7. 保存和加载模型时指定设备:
   当你保存或加载模型时，确保模型在正确的设备上。

   # 保存模型
   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，例如：

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

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

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