PyTorch指定GPU的方法

分享好文章，本文出自简书作者-几时见得清梦

1. 改变系统变量使得仅目标显卡可见

1. export CUDA_VISIBLE_DEVICES=0 #这里是要使用的GPU编号，正常的话是从0开始
2. 在程序开头设置os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
3. CUDA_VISIBLE_DEVICES=0,1,2,3 python xxx.py来设置该程序可见的gpu

2. 使用torch.cuda接口

在生成网络对象之前执行：torch.cuda.set_device(0)

3. 使用pytorch的并行GPU接口

net = torch.nn.DataParallel(model, device_ids=[0])

4. 初始化模型时

net = Net.cuda(0)

5. 在pycharm中调试时如何指定GPU

1. 直接在程序中设置os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'

2. 打开Run/Debug Configurations，在Environment variables右边点击…处。

    

   

3. pycharm设置gpu

4. 点击下图中1处的’+’，会出现2处，增加CUDA_VISIBLE_DEVICES，再点击3处，填入0,1,2,3。（下图中第一行已经增加了，只不过为了说明，因而这样表示）

    

   

   pycharm设置gpu

5. 在pycharm中正常运行程序即可

作者：几时见得清梦
链接：https://www.jianshu.com/p/a014016723d8
来源：简书
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。