pytorch如何查看tensor和model在哪个GPU上以及指定GPU设备

本文介绍了如何在PyTorch中查看和操作模型及张量在GPU上的位置。通过`.cuda()`和`.cpu()`方法实现数据和模型在GPU与CPU之间的转移。同时,提供了两种常见方法将模型和张量加载到GPU,以及如何指定使用特定GPU。此外,还展示了检查GPU数量的代码。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

1. 查看tensor所在的设备:
data = data.cuda()#将数据转移到gpu上
 
print(data.device)  # 输出:cuda:0
 
data = data.cpu()#将数据转移到cpu上
 
print(data.device)  # 输出:cpu
2. 查看model所在的设备
model = model.cuda()#将模型转移到gpu上
 
print(next(model.parameters()).device)  # 输出:cuda:0
 
model = model.cpu()#将模型转移到cpu上
 
print(next(model.parameters()).device)  # 输出:cpu
3. Pytorch中将模型和张量加载到GPU的常用方法有两种。

方式1:

# 如果GPU可用,将模型和张量加载到GPU上
if torch.cuda.is_available():
    model = model.cuda()
    x = x.cuda()
    y = y.cuda()

方式2:

# 分配到的GPU或CPU
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将模型加到GPU
model=model.to(device)
# 将张量加到GPU
x=x.to(device)
y=y.to(device)
4. 指定GPU代码
# 代码1:
torch.cuda.set_device(1)

# 代码2:
device = torch.device("cuda:1")

# 代码3:(官方推荐使用),
os.environ["CUDA_VISIBLE_DEVICES"] = '1'

(如果你想同时调用两块GPU的话)
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'

参考链接:PyTorch 中 选择指定的 GPU

注意需要将指定GPU代码放在程序段最开始的部位,如下图所示:
gpu

5.查看gpu个数
torch.cuda.device_count()
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值