5.6.0 头文件
import torch
from torch import nn
5.6.1获取GPU
# 获取指定GPU,如果指定编号的GPU存在,则返回gpu(i),否则返回cpu
def try_gpu(i=0):
if torch.cuda.device_count() >= i + 1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
# 获取所有GPU,如果没有GPU,则返回CPU
def try_all_gpus():
devices = [torch.device(f'cuda:{i}')
for i in range(torch.cuda.device_count())]
return devices if devices else [torch.device('cpu')]
print(try_gpu())
# 输出:
# cuda:0
print(try_gpu(10))
# 输出:
# cpu
print(try_all_gpus())
# 输出:
# [device(type='cuda', index=0)]
5.6.2 在CPU或GPU上创建张量
# 在CPU上创建张量
cpu_x = torch.tensor([1,2,3])
# 查看张量在什么设备上
print(cpu_x.device)
# 输出:
# cpu
# 在GPU上创建张量
gpu_x = torch.ones(2, 3, device=try_gpu())
# 查看张量在什么设备上
print(gpu_x.device)
# 输出:
# cuda:0
# 在指定GPU上创建张量
gpu_x = torch.ones(2, 3, device=try_gpu(0))
# 查看张量在什么设备上
print(gpu_x.device)
# 输出:
# cuda:0
5.6.3 用GPU来计算网络模型
# 定义一个神经网络模型,并将模型参数放在GPU上
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
# 当输入为同一个GPU上的张量时,模型将在同一个GPU上计算结果
print(net(gpu_x))
# 输出:
# tensor([[-0.1049],
# [-0.1049]], device='cuda:0', grad_fn=<AddmmBackward0>)
# 查看模型参数位于哪个设备上
print(net[0].weight.data.device)
# 输出:
# cuda:0
5.6.4 其他
# 查看cuda是否可用
print(torch.cuda.is_available())
# 查看有几个GPU
print(torch.cuda.device_count())
本小节完整代码如下
import torch
from torch import nn
# 查看cuda是否可用
print(torch.cuda.is_available())
# 查看有几个GPU
print(torch.cuda.device_count())
# ------------------------------获取GPU------------------------------------
# 获取指定GPU,如果指定编号的GPU存在,则返回gpu(i),否则返回cpu
def try_gpu(i=0):
if torch.cuda.device_count() >= i + 1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
# 获取所有GPU,如果没有GPU,则返回CPU
def try_all_gpus():
devices = [torch.device(f'cuda:{i}')
for i in range(torch.cuda.device_count())]
return devices if devices else [torch.device('cpu')]
print(try_gpu())
# 输出:
# cuda:0
print(try_gpu(10))
# 输出:
# cpu
print(try_all_gpus())
# 输出:
# [device(type='cuda', index=0)]
# ------------------------------在CPU或GPU上创建张量------------------------------------
# 在CPU上创建张量
cpu_x = torch.tensor([1,2,3])
# 查看张量在什么设备上
print(cpu_x.device)
# 输出:
# cpu
# 在GPU上创建张量
gpu_x = torch.ones(2, 3, device=try_gpu())
# 查看张量在什么设备上
print(gpu_x.device)
# 输出:
# cuda:0
# 在指定GPU上创建张量
gpu_x = torch.ones(2, 3, device=try_gpu(0))
# 查看张量在什么设备上
print(gpu_x.device)
# 输出:
# cuda:0
# ------------------------------用GPU来计算网络模型------------------------------------
# 定义一个神经网络模型,并将模型参数放在GPU上
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
# 当输入为同一个GPU上的张量时,模型将在同一个GPU上计算结果
print(net(gpu_x))
# 输出:
# tensor([[-0.1049],
# [-0.1049]], device='cuda:0', grad_fn=<AddmmBackward0>)
# 查看模型参数位于哪个设备上
print(net[0].weight.data.device)
# 输出:
# cuda:0
# ------------------------------其他------------------------------------
# 查看cuda是否可用
print(torch.cuda.is_available())
# 查看有几个GPU
print(torch.cuda.device_count())