温故而知新,可以为师矣!
一、参考资料
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
-
直接终端中设定:
CUDA_VISIBLE_DEVICES=1
-
python代码中设定:
import os os.environ['CUDA_VISIBLE_DEVICE']='1'
-
使用
set_device()
函数:import torch # Pytoch中的`in-place operation`原地操作符 torch.cuda.set_device(id)
4. 相关经验
- Tensor 和 Numpy都是矩阵,区别是前者可以在GPU上运行,后者只能在CPU上;
- Tensor可以直接通过print显示数据类型,而Numpy不可以;