参考: [开发技巧]·PyTorch如何使用GPU加速(CPU与GPU数据的相互转换)
安装可以在官方参考代码,最好使用conda https://pytorch.org/get-started/locally/
conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
检测是否可以使用GPU,使用一个全局变量use_gpu,便于后面操作使用
use_gpu = torch.cuda.is_available()
可以使用GPU,use_gpu的值为True,否则为False。当可以使用GPU,我们不想使用,可以直接赋值use_gpu = False
我们在进行转换时,需要把数据,网络,与损失函数转换到GPU上
1.构建网络时,把网络,与损失函数转换到GPU上
model = get_model()
loss_f = t.nn.CrossEntropyLoss()
if(use_gpu):
model = model.cuda()
loss_f = loss_f.cuda()
2.训练网络时,把数据转换到GPU上
if (use_gpu):
x,y = x.cuda(),y.cuda()
3.取出数据是,需要从GPU准换到CPU上进行操作
if(use_gpu):
loss = loss.cpu()
acc = acc.cpu()
未完待续