深度学习之python使用指定gpu运行代码

一、前提

1、在命令行使用nvidia-smi查看gpu设备情况,当存在空闲设备时才能用,否则会出现运行内存不够出错的情况(具体参考文章GPU之nvidia-smi命令详解);

2、安装好cuda和cudcnn(具体步骤请参考:①windows:CUDA安装教程(超详细)),②linux:linux安装CUDA+cuDNN)

在命令行输入nvcc -V或者nvcc --version检查是否安装成功:

检查CUDA是否安装成功

二、命令行直接指定

运行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.eval()
model.to('cuda')

参考自python指定gpu运行

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
心中带点小风骚的头像心中带点小风骚普通用户
上一篇 2023年8月3日
下一篇 2023年8月3日

相关推荐