PyTorch之Tensor.to(device)和model.to(device)详解

温故而知新，可以为师矣！

一、参考资料

pytorch中Tensor.to(device)和model.to(device)的区别及说明
Pytorch to（device）
Pytorch的to(device)用法

二、PyTorch之to(device)

在PyTorch中，可以使用to()方法将Tensor或模型移动到指定的设备上。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
model.to(device)

以上两行代码放在读取数据之前。

mytensor = my_tensor.to(device)

这行代码的意思是，将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去，之后的运算都在GPU上进行。

1. Tensor.to(device)

功能：将Tensor移动到指定的设备上。

可以使用以下代码将Tensor移动到GPU上：

import torch


x = torch.randn(3, 3)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
x = x.to(device)

在这个示例中，首先创建一个3x3的Tensor x，然后使用 torch.device() 函数创建一个设备对象device。如果GPU可用，则将device设置为’cuda:0’，否则将device设置为’cpu’。最后，使用x.to(device) 将Tensor x移动到指定的设备上。

import torch


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

data_transform = {
    "train": transforms.Compose([transforms.Resize((224, 224)),
                                 transforms.ToTensor(),
                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.255])]),
    "test": transforms.Compose([transforms.Resize((224, 224)),
                                transforms.ToTensor(),
                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.255])])}

# 训练集
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=data_transform['train'])
trainloader = DataLoader(trainset, batch_size=16, shuffle=True)

for i, data in enumerate(trainloader):
    # 获取输入
    inputs, labels = data
    inputs = inputs.to(device)  
    labels = labels.to(device)

上述实例中，将 inputs和labels数据移动到指定的设备上。

2. model.to(device)

功能：将模型移动到指定的设备上。

可以使用以下代码将模型移动到GPU上：

import torch
import torch.nn as nn


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

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

    
model = Net()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(device)

在这个示例中，我们首先定义了一个简单的神经网络模型Net，并使用 net.to(device) 将模型移动到指定的设备上。其中，device=torch.device("cpu")表示使用cpu，而device=torch.device("cuda")则表示使用GPU。

如果有多个GPU，使用以下方法：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
model = Model()
 
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model，device_ids=[0,1,2])
  
model.to(device)

将由GPU保存的模型加载到GPU上。

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)

将由GPU保存的模型加载到CPU上。将torch.load()函数中的map_location参数设置为torch.device('cpu')。

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
model.to(device)

将由CPU保存的模型加载到GPU上。将torch.load()函数中的map_location参数设置为torch.device('cuda')。

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)

3. PyTorch中使用指定的GPU

1. 直接终端中设定：

   CUDA_VISIBLE_DEVICES=1
   

2. python代码中设定：

   import os
   
   
   os.environ['CUDA_VISIBLE_DEVICE']='1'
   

3. 使用set_device()函数：

   import torch
   
   
   # Pytoch中的`in-place operation`原地操作符
   torch.cuda.set_device(id)
   

4. 相关经验

• Tensor 和 Numpy都是矩阵，区别是前者可以在GPU上运行，后者只能在CPU上；
• Tensor可以直接通过print显示数据类型，而Numpy不可以；