jupyter中如何使用多GPU训练模型[pytorch]

jupyter中如何使用多GPU训练模型[pytorch]

训练模型过程中，可能因为显存不够等原因，我们需要使用多个卡进行运行程序。
默认的一句话，这个是需要的

device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
model = model.to(device)

* 在终端中运行时候

在程序开始的位置插入一句话

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" 

这样在之后的模型训练的时候，自动就会使用多卡进行训练了

* 在jupyter中运行的时候

发现在使用jupyter中运行的时候，插入上述代码没有效果
可以使用下面 的方法

model = nn.DataParallel(model, device_ids=[0,1,2,3])

这样就可以指定在0,1,2,3四张卡上面运行了。

参考资料：
https://blog.csdn.net/qq_43307074/article/details/127662174