运行improved-diffusion训练CelebA时遇TCPStore格式字符串错误求助

问题背景 #

• 基于OpenAI的improved-diffusion仓库，使用CelebA人脸数据集训练图像修复类扩散模型
• 执行训练命令：python scripts/image_train.py --data_dir "D:\WeiTsung\improved_diffusion_main\DDPM_train\img_align_celeba"
• 设备配置：2块GeForce RTX4070 12GB GPU，torch.cuda.device_count()返回2，torch.cuda.is_available()返回True，GPU可正常访问
• 已尝试更换PyTorch和CUDA版本，问题未解决

错误日志 #

Traceback (most recent call last):
  File "D:\WeiTsung\improved_diffusion_main\DDPM_train\scripts\image_train.py", line 87, in <module>
    main()
  File "D:\WeiTsung\improved_diffusion_main\DDPM_train\scripts\image_train.py", line 26, in main
    dist_util.setup_dist()
  File "D:\WeiTsung\improved_diffusion_main\DDPM_train\improved_diffusion\dist_util.py", line 51, in setup_dist
    dist.init_process_group(backend=backend, init_method="env://")
  File "C:\Users\user\anaconda3\envs\DDPM_training_v2\Lib\site-packages\torch\distributed\c10d_logger.py", line 74, in wrapper
    func_return = func(*args, **kwargs)
                  ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\user\anaconda3\envs\DDPM_training_v2\Lib\site-packages\torch\distributed\distributed_c10d.py", line 1145, in init_process_group
    store, rank, world_size = next(rendezvous_iterator)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\user\anaconda3\envs\DDPM_training_v2\Lib\site-packages\torch\distributed\rendezvous.py", line 247, in _env_rendezvous_handler
    store = _create_c10d_store(master_addr, master_port, rank, world_size, timeout)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\user\anaconda3\envs\DDPM_training_v2\Lib\site-packages\torch\distributed\rendezvous.py", line 178, in _create_c10d_store
    return TCPStore(
           ^^^^^^^^^
RuntimeError: unmatched '}' in format string

TCPStore传入参数 #

hostname=140.117.172.120
port=2021（每次运行端口值会变化）
world_size=1
start_daemon=True
timeout=0:30:00

解决方案 #

1. 清除并重置分布式环境变量 #

错误根源是环境变量中存在未闭合的{}字符，导致TCPStore解析格式字符串失败。运行训练命令前先手动设置正确的环境变量：

# Windows系统
set MASTER_ADDR=127.0.0.1
set MASTER_PORT=29500

# Linux/macOS系统
export MASTER_ADDR=127.0.0.1
export MASTER_PORT=29500

2. 显式指定分布式启动参数 #

使用PyTorch的分布式启动器，直接指定GPU数量，避免依赖自动环境变量解析：

python -m torch.distributed.launch --nproc_per_node=2 scripts/image_train.py --data_dir "D:\WeiTsung\improved_diffusion_main\DDPM_train\img_align_celeba"

其中--nproc_per_node=2指定使用2块GPU参与训练。

3. 修改分布式初始化逻辑 #

如果上述方法无效，手动修改improved_diffusion/dist_util.py中的setup_dist函数，显式定义分布式参数：

def setup_dist():
    if dist.is_initialized():
        return
    # 强制指定可见GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
    dist.init_process_group(
        backend="nccl",  # GPU分布式训练推荐后端
        init_method="tcp://127.0.0.1:29500",
        rank=int(os.environ.get("RANK", 0)),
        world_size=int(os.environ.get("WORLD_SIZE", 2))
    )

4. 检查路径特殊字符 #

确保数据路径、项目路径中不包含{}、%等会触发格式解析的特殊字符，避免干扰环境变量或参数解析。

内容的提问来源于stack exchange，提问作者ChuckLin