功能:将Tensor移动到指定的设备上。
以下代码将Tensor移动到GPU上:
device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)
a = tensor.to(torch.float64)
.
a = tensor.to('cuda:0')
.
cuda0 = torch.device('cuda:0')
.
b = tensor.to(cuda0)
.
c = tensor.to('cuda:0',torch.float64)
.
other = torch.randn((), dtype=torch.float64, device=cuda0)
.
d = tensor.to(other, non_blocking=True)
.
功能:将模型移动到指定的设备上。
使用以下代码将模型移动到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)
如果有多个GPU,使用以下方法:
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)
参考:PyTorch之Tensor.to(device)和model.to(device)
图参考:pytorch中.to(device) 和.cuda()的区别
官方文档:CUDA SEMANTICS
with torch.cuda.device(1):
# allocates a tensor on GPU 1
a = torch.tensor([1., 2.], device=cuda)
# transfers a tensor from CPU to GPU 1
b = torch.tensor([1., 2.]).cuda()
# a.device and b.device are device(type='cuda', index=1)
# You can also use ``Tensor.to`` to transfer a tensor:
b2 = torch.tensor([1., 2.]).to(device=cuda)
# b.device and b2.device are device(type='cuda', index=1)
import torch
print('CUDA版本:',torch.version.cuda)
print('Pytorch版本:',torch.__version__)
print('显卡是否可用:','可用' if(torch.cuda.is_available()) else '不可用')
print('显卡数量:',torch.cuda.device_count())
print('当前显卡的CUDA算力:',torch.cuda.get_device_capability(0))
print('当前显卡型号:',torch.cuda.get_device_name(0))
>>>
CUDA版本: 11.7
Pytorch版本: 1.13.1
显卡是否可用: 可用
显卡数量: 1
当前显卡的CUDA算力: (8, 6)
当前显卡型号: NVIDIA GeForce RTX 3060 Laptop GPU
参考:https://blog.csdn.net/weixin_43845386/article/details/131723010