原先的代码可能如下:
device = torch.device("cuda:0" if torch.cuda.is_available() and not args.no_cuda else "cpu")
model = model.to(device)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
上面这段代码会指定序号为0的GPU作为初始加载的GPU,同时如果有多GPU的情况下,会使用该服务器上所有的GPU运行程序。假定服务器有8块GPU,序号就是0-7,该代码就会全部使用0-7的GPU。但是如果服务器上假设序号为0-3的GPU被人占用了,那怎么办?直接使用上述代码就会OOM,所以对上述代码修改,就是下面的代码:
device = torch.device("cuda:{}".format(4) if torch.cuda.is_available() and not args.no_cuda else "cpu")
model = model.to(device)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model, device_ids=[4,5,6,7])
这段代码将初始加载的GPU改为序号为4的GPU,并且在多GPU的情况下,选取了序号为4-7的GPU来运行代码,这样子在序号为0-3的GPU被占用时就不会和别人冲突了。