Pytorch在多GPU下选择特定序号的GPU运行程序

原先的代码可能如下：

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被占用时就不会和别人冲突了。