pytorch多卡训练DistributedDataParallel 显卡开始不是0

先说两个试出来的方法

1.每个local_rank都加一个值

比如默认用两个显卡，他会给0，1你想用1，2就都加1

#比如
torch.cuda.set_device(args.local_rank+2)

2.在终端命令前加上CUDA_VISIBLE_DEVICES=2,3

我尝试在代码里面加入不知为什么没有任何效果，还是要在终端上家上面的那句话

CUDA_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 modeltest.py

具体情况和代码

今天我尝试了很久多卡训练，当然是不太想使用DataParallel，因此还是反复尝试，一直报错，最后探索出了一个可行方法，我也不是很清楚细节。

在运行时貌似我所有的print都会输出两次。

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=os.getenv('LOCAL_RANK', -1), type=int)
args = parser.parse_args()
print(args.local_rank)
if args.local_rank != -1:
    torch.cuda.set_device(args.local_rank+2)
    device=torch.device("cuda", args.local_rank+2)
    torch.distributed.init_process_group(backend="nccl", init_method='env://')

我要说明上面的local_rank的问题，这里面的localrank就应该是cuda的序号，而由于在终端指定了两个，它默认会输出0，1，然而我这里是想用2，3的卡。所以我在所有的部分都加了一个2

net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank+2], output_device=args.local_rank+2)

最后还是有效果的。

当然model 和 数据也要.to(device)

net.to(device)

#下面这个是图片和标签，大家自己注意加入的时机
img.to(device)
target.to(device)

我也尝试了在代码中直接指定os.environ可见devices但没有效果

有任何详细的解释大家可以打在评论区，我也对此很懵hh，求大佬