一、前提
1、在命令行使用nvidia-smi
查看gpu设备情况,当存在空闲设备时才能用,否则会出现运行内存不够出错的情况(具体参考文章GPU之nvidia-smi命令详解);
在命令行输入nvcc -V
或者nvcc --version
检查是否安装成功:
![检查CUDA是否安装成功](https://i-blog.csdnimg.cn/blog_migrate/f43ec060e978164c41db0f69ec6f800a.png)
二、命令行直接指定
运行python文件时,在前面加上 CUDA_VISIBLE_DEVICES=xxx,xxx是你要指定的gpu序号。如CUDA_VISIBLE_DEVICES=0,2 python extract_masks.py
表示指定了0、2号两个gpu。注意,序号是从0开始排的,所以如果你只有一个gpu,那么就是CUDA_VISIBLE_DEVICES=0 python extract_masks.py
。
三、在python文件中指定
在需要运行的python文件最开始加上两行如下,等号后面的数字同上所述。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2,3"
四、如何使用
在python文件中,定义需要加速的模型之后,加上:
model = ......
model.train(True) or model.train(False)
model.to('cuda')
参考自python指定gpu运行