model = Model() if torch.cuda.is_available(): model = model.cuda() 123
上面代码直接使用model.cuda(), PyTorch默认使用从0开始的GPU, 有如下两种方法来指定需要使用的GPU:
直接在终端中设定：CUDA_VISIBLE_DEVICES=1 python main.pypython代码中设定：import os os.environ["CUDA_VISIBLE_DEVICES"] = "1" model = Model() if torch.cuda.is_available():model = model.cuda() #使用第一个GPU 12345 2. PyTorch使用指定GPU训练
device_ids = [3, 4, 6, 7] model = Module() if torch.cuda.is_available(): module = torch.nn.DataParallel(model, device_ids=device_ids) # 声明所有可用设备 model = model.cuda(device=device_ids[0]) # 模型放在主设备 images = images.cuda(device=device_ids[0]) # 训练数据放在主设备 labels = labels.cuda(device=device_ids[0]) 12345678
使用多GPU训练，model = nn.DataParallel(model)
注意训练/测试过程中 inputs和labels均需加载到GPU中
