如何在GPU上运行pytorch程序(正确方法)

用GPU跑pytorch程序就3点:

1.申明用GPU

2.把你的model放到GPU上

3.把数据和标签放到GPU上

详细步骤:

1.申明:

device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

如果有多个GPU,你可以选择具体的GPU进行使用。

import os 

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]='1'

选第几块就是填对应的序号就行。

比如这里我是用的第二块,但是序号为1:

可以看出来,我用的是1.

可以通过watch -n 3 nvidia-smi来查看gpu的使用状况(这个语句是3s刷新一次)

如果多个并行比较复杂。请自行百度。

2.将模型放到GPU上

在创建完网络 或者引用网络之后,我们需要实体化我们的网络。

直接在后面加一句话就可以

net= Net ()
net.to(device)

3.把数据放到GPU上

inputs,labels=data

inputs, labels =  data[0].to(device), data[1].to(device)

或者

inputs, labels= inputs.to(device),labels.to(device)

 

这样就可以跑网络了。

 

  • 31
    点赞
  • 103
    收藏
    觉得还不错? 一键收藏
  • 4
    评论
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值