https://blog.csdn.net/wacebb/article/details/114919245
用GPU跑pytorch程序就3点:
1.申明用GPU
2.把你的model放到GPU上
3.把数据和标签放到GPU上
一、申明:
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)
这样就可以跑网络了。