利用yolov5做项目期间,由于要同时训练两个项目代码,发现一个之前没有发现的问题:GPU调用默认从0开始。
举个例子:train.py中device调用代码
parser.add_argument('--device', default='4,5,6,7', help='device id (i.e. 0 or 0,1 or cpu)')
比如我这里在这里调用 device = 4,5,6,7 。按理来讲,实验室服务器8块GPU,应该调用后面四个才对。然而实际上它默认从0开始,调用了0,1,2,3四块GPU。控制变量,修改代码为
parser.add_argument('--device', default='4,5', help='device id (i.e. 0 or 0,1 or cpu)')
这时服务期显示调用了0,1两个GPU。所以问题如下:
在小师姐指导下,在import后添加一行代码
os.environ["CUDA_VISIBLE_DEVICES"] = "4,5,6,7"
完美解决问题。