0
点赞
收藏
分享

微信扫一扫

LTE和5G基站规划原则

年迈的代码机器 2024-06-23 阅读 32

18GPU

import torch
from torch import nn

#使用CPU
torch.device('cpu')

#使用GPU
#cuda:0和cuda是等价的
#是一个上下文管理器,用于在其上下文范围内临时设置默认CUDA设备
torch.cuda.device('cuda')

#如果有多个GPU,torch.device(f'cuda:{i}'),表示使用第i块GPU(从0开始)
#是一个设备对象,用于在创建或移动张量时明确指定目标设备
torch.device('cuda:1')

# 查询可用gpu的数量
print(torch.cuda.device_count())

def try_gpu(i = 0):
    """如果存在,则返回gpu(i),否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():
    """返回所有可用的GPU,如果没有GPU,则返回[cpu(),]"""
    devices = [torch.device(f'cuda:{i}')
               for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]
# 如果 devices 列表不为空,返回 devices 列表。
# 如果 devices 列表为空,返回一个包含单个CPU设备的列表 [torch.device('cpu')]。

print(try_gpu())
print(try_gpu(10))
print(try_all_gpus())

#张量与GPU

x = torch.tensor([1, 2, 3])
print(x.device)
# cpu
# 默认情况下,张量是在CPU上创建的。

#张量存储在GPU上
Y = torch.ones(2, 3, device=torch.device('cuda'))
X = torch.ones(2, 3, device=try_gpu())
print(Y)
print(X)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]], device='cuda:0')

# 假设至少有两个GPU
# Y = torch.rand(2, 3, device=try_gpu(1))
# print(Y)
# tensor([[0.4860, 0.1285, 0.0440],
#         [0.9743, 0.4159, 0.9979]], device='cuda:1')

#复制
# Z = X.cuda(1)
# print(X)
# print(Z)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]], device='cuda:0')
# tensor([[1., 1., 1.],
#         [1., 1., 1.]], device='cuda:1')

#相加
print(X + Y)
# tensor([[2., 2., 2.],
#         [2., 2., 2.]], device='cuda:0')

# 假设变量Z已经存在于第二个GPU上。 如果我们还是调用Z.cuda(1)
# 不会复制并分配新内存
# Z.cuda(1) is Z
# True

# 神经网络与GPU
# 神经网络模型可以指定设备
net = nn.Sequential(nn.Linear(3, 1)) # 在cup上建立
net = net.to(device=try_gpu()) # 将模型参数放在GPU上

print(net(X))
# tensor([[-1.4124],
#         [-1.4124]], device='cuda:0', grad_fn=<AddmmBackward>)

#查看模型参数在哪里
print(net[0].weight.data.device)
# cuda:0
举报

相关推荐

0 条评论