pytorch如何查看tensor和model在哪个GPU上以及指定GPU设备

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()