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

2024-02-10 09:45| 来源: 网络整理| 查看: 265

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

一、参考资料

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不可以；

【本文地址】

公司简介

联系我们