pytorch nn.DataParallel指定多GPU运行

有两种方法可以指定gpu数目运行，第一种进入ubuntu虚拟环境后在终端使用EXPORT CUDA_VISIBLE_DEVICES=[你想要训练的GPU序号]，可以通过nvidia-smi来查看空闲的GPU序号，同时如果batch_size设置过大，会导致发生cuda out of memory的报错，这个时候通过新建一个窗口输入watch -n 5 nvidia-smi，即每隔5s刷新GPU使用状态，来观察是否有被调用GPU的行为发生，另一种方法即是在程序中模型加载前使用import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"指定你想要训练GPU的序号即可。