Pytorch模型model&data.to(device) | .cuda | .cpu()

模型model或数据data放到cpu或gpu上

模型和数据需要在同一个设备上,才能正常运行:

  1. model和data都在cpu上
  2. model和data都在gpu上
model = Model()
input = dataloader()
output = model(input)

1. to.(device)

device自己定义好,可以是cpu或者gpu

device = 'gpu'
model = model.to(device)
data = data.to(device)
# 
device = 'cpu'
model = model.to(device)
data = data.to(device)

在这里插入图片描述

2. .cuda()

model = model.cuda()
data = data.cuda()

在这里插入图片描述

3. .cpu()

model = model.cpu()
data = data.cpu()

4. 检查

  1. 通过判断模型model的参数是否在cuda上来判定模型是否在gpu上。
print('Is model on gpu: ', next(model.parameters()).is_cuda)

输出若是True,则model在gpu上;若是False,则model在cpu上。

  1. 输出数据data的device字段。
print('data device: ', data.device)

输出gpu则在gpu上,输出cpu则在cpu上。

————————————————
https://blog.csdn.net/iLOVEJohnny/article/details/106021547
https://wenku.baidu.com/view/99bec6ef8aeb172ded630b1c59eef8c75fbf95ed.html

猜你喜欢

转载自blog.csdn.net/LemonShy2019/article/details/126156250