分布式训练environment variable RANK expected, but not set

(83条消息) 报错: environment variable RANK expected, but not set__湘江夜话_的博客-CSDN博客------------------------------------------------------------

2. 上述方式不方便调试，可将

def setup_dist():
    """
    Setup a distributed process group.
    Return (bool, int) which indicates whether to use distrubted training
        and which GPU rank is used for current device.
    """
    if th.cuda.is_available():
        # os.environ['LOCAL_RANK'] = '0'
        th.cuda.set_device(int(os.environ["LOCAL_RANK"]))
        dist.init_process_group(backend="nccl", init_method="env://")
        return True, int(os.environ["RANK"])
    else:
        return False, 0
# python -m torch.distributed.launch test.py  # xxx.py为要执行的脚本文件名

该函数注释掉。