Torch之CUDA使用指定显卡

1. 改变系统环境变量仅使目标显卡,编辑 .bashrc文件，添加系统变量

export CUDA_VISIBLE_DEVICES=0 #这里是要使用的GPU编号，正常的话是从0开始

2. 使用torch.cuda接口函数

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

3. 运行文件之前指定显卡

CUDA_VISIBLE_DEVICES=7 python retinanet/main.py

4. 使用多torch的并行GPU接口

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

5. 初始化网络模型时

net = Net.cuda(0)

注：

参考博客：https://blog.csdn.net/weixin_40712763/article/details/82846509