使用Pytorch进行多GPU训练数据

今天尝试使用多显卡训练PyTorch代码。

不过训练时报错，查找代码发现此前曾设定单一GPU训练，如下。

将其注释掉，再次运行，又发现报错。

查找资料发现训练模型所使用的显卡必须要包括主卡，而我在设置训练GPU时，默认0卡为主卡，但却没有使用，所以会报错。

gpu_ids = '1, 2'

如果要将其他卡设置为主卡，例如我们这里将卡1设为主卡，卡2并行训练，则需要在训练脚本中添加如下代码，表明让程序可视卡1、卡2。此时实际显卡编号“1, 2”的索引分别为“0, 1”，所以我们需要更改gpu_ids=‘0, 1’，表明使用运算显卡编号“0, 1”进行训练。

os.environ["CUDA_VISIBLE_DEVICES"] = "1, 2"

至此，代码已经修改完毕，可以正常多GPU训练。