PyTorch查看模型和数据是否在GPU上

模型和数据可以在CPU和GPU上来回迁移，怎么判断模型和数据在哪里呢？

?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

import torch import torch.nn as nn   # ----------- 判断模型是在CPU还是GPU上 ----------------------   model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True) print(next(model.parameters()).device)  # 输出：cpu   model = model.cuda() print(next(model.parameters()).device)  # 输出：cuda:0   model = model.cpu() print(next(model.parameters()).device)  # 输出：cpu   # ----------- 判断数据是在CPU还是GPU上 ----------------------   data = torch.ones([2, 3]) print(data.device)  # 输出：cpu   data = data.cuda() print(data.device)  # 输出：cuda:0   data = data.cpu() print(data.device)  # 输出：cpu

此外，用.is_cuda也可以判断模型和数据是否在GPU上，例如： data.is_cuda 

模型和数据在GPU和CPU之间进行迁移的方法，更具体的可以参考：把数据和模型迁移到GPU: .cuda()方法和.to(device)方法