pytorch 中判断和指定模型和数据在GPU或CPU上

有时候,需要查看模型和数据是在GPU上,还是在CPU上;或者需要将模型和数据放在指定的GPU或者CPU上,那该怎么做呢?

  • 1. 判断模型是在GPU上还是CPU上
model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).device)  
  • 2.判断数据是在GPU上还是CPU上
data = torch.ones([2, 3])
print(data.device) 
  • 3.指定模型是在GPU上还是CPU上
model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
# 指定模型在GPU上
model = model.cuda()
# 指定模型在CPU上
model = model.cpu()
  • 4.指定数据是在GPU上还是在CPU上
data = torch.ones([2, 3])
# 指定数据在GPU上 
data = data.cuda()
# 指定数据在CPU上 
data = data.cpu()

参考:https://www.cnblogs.com/picassooo/p/13736843.html

猜你喜欢

转载自blog.csdn.net/m0_46483236/article/details/123942225