深度学习中gpu的写法
文章目录
- 第一类 在命令行中指定
- 第一种写法
- 第二种写法
- 第二类 在代码中指定
- 第一种写法
本文讲解的是在深度学习的代码中常见的几种设置模型运行在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().