1. 查看tensor所在的设备:
data = data.cuda()#将数据转移到gpu上 print(data.device)# 输出:cuda:0 data = data.cpu()#将数据转移到cpu上 print(data.device)# 输出:cpu
2. 查看model所在的设备
model = model.cuda()#将模型转移到gpu上 print(next(model.parameters()).device)# 输出:cuda:0 model = model.cpu()#将模型转移到cpu上 print(next(model.parameters()).device)# 输出:cpu
3. Pytorch中将模型和张量加载到GPU的常用方法有两种。
方式1:
# 如果GPU可用,将模型和张量加载到GPU上if torch.cuda.is_available():model = model.cuda()x = x.cuda()y = y.cuda()
方式2:
# 分配到的GPU或CPUdevice=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# 将模型加到GPUmodel=model.to(device)# 将张量加到GPUx=x.to(device)y=y.to(device)
4. 指定GPU代码
# 代码1:torch.cuda.set_device(1)# 代码2:device = torch.device("cuda:1")# 代码3:(官方推荐使用),os.environ["CUDA_VISIBLE_DEVICES"] = '1'(如果你想同时调用两块GPU的话)os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
参考链接:PyTorch 中 选择指定的 GPU
注意需要将指定GPU代码放在程序段最开始的部位,如下图所示:
5.查看gpu个数
torch.cuda.device_count()