深度学习中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().