5.6 GPU与张量，GPU与网络模型

 

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())