模型model或数据data放到cpu或gpu上
模型和数据需要在同一个设备上,才能正常运行:
- model和data都在cpu上
- 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. 检查
- 通过判断模型model的参数是否在cuda上来判定模型是否在gpu上。
print('Is model on gpu: ', next(model.parameters()).is_cuda)
输出若是True,则model在gpu上;若是False,则model在cpu上。
- 输出数据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