文章目录
- 第一类 在命令行中指定
- 第一种写法
- 第二种写法
- 第二类 在代码中指定
- 第一种写法
本文讲解的是在深度学习的代码中常见的几种设置模型运行在gpu上的写法,主要分为两大类:
- 第一类是:在命令行中指定gpu卡号
- 第二类是:在代码中指定gpu卡号
第一类 在命令行中指定
第一种写法
parser.add_argument("--gpu", type=int, default=0)
device = torch.device(args.gpu) # 或者是 device = f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu'
采用此方法定义gpu卡号,那么在使用的时候,比如我们要把一个模型放到gpu上,可以这样写model = model.to(device)
.
第二种写法
parser.add_argument("--gpu", type=int, default=0)
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
使用方法和上面不同,此方法使用方式为model = model.cuda()
.
第二类 在代码中指定
第一种写法
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
这句话不管放在哪里都行,比如if __name = "__main__
或者import的时候
。使用方法也很简单,model = model.cuda()
.