1. cpu()
对于模型的cpu()函数,比如 model.cpu(),只是把 parameters 和 buffers 进行转移
2. cuda(device=None)
在PyTorch中,算法优先使用CPU,不会自动使用GPU,因此需要在程序中显示指定。如果模型要在gpu上进行训练,需要再将模型参数放到优化器之前调用cuda() 将其加载到GPU。
- 在网络模型的训练过程中,模型,数据(input和gt)和损失函数都需要转移到cuda
- 使用 torch.cuda.is_available() 判断gpu是否可用
- 除了 cuda(device) 函数外,还可以使用 to(device) 来进行设备转换。其中 to(device) 中的参数是必须的,cuda()中的参数是可选的。对于单机单卡(单GPU)的模型训练,代码如下:
model = Model()
# 第一种定义 device的方式
device = torch.device("cuda:1")
# 第二种定义 device的方式
device = torch.dev