基于pytorch训练卷积神经网络,在多gpu环境下并行训练,会有一个主显卡,相比于其他显卡,它还需统计其他显卡发送的梯度等参数进行下一步计算,因此需要更大的显存。默认情况下,会以0号显卡为主显卡。若0号显卡被其他程序占部分内存,为主显卡可能会因为内存不足oom,可以将指定其他显卡为主显卡,在程序运行命令前加:CUDA_VISIBLE_DEVICES=3,0,1,2 就会指定3号为主显卡,且在4块GPU上并行训练。例:CUDA_VISIBLE_DEVICES=3,0,1,2 python main.py