基于Weights & Biases Sweep的多进程并行配置与实操问询

我正在用Weights and Biases（W&B）做超参数调优，想要通过多进程实现最大程度的实验并行，保证每组超参数只评估一次，多实验同时运行不重复。以下是我的训练与评估简化脚本：

import random
import numpy as np
import wandb
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Tuple, Dict

def train_one_epoch(epoch: int, lr: float, bs: int) -> Tuple[float, float]:
    """
    Simulate training for one epoch.
    """
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch: int) -> Tuple[float, float]:
    """
    Simulate evaluation for one epoch.
    """
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

config: Dict[str, float] = {"lr": 0.0001, "bs": 16, "epochs": 5}

def run_epoch(epoch: int, lr: float, bs: int) -> Tuple[int, float, float, float, float]:
    """
    Run training and evaluation for one epoch.
    """
    train_acc, train_loss = train_one_epoch(epoch, lr, bs)
    val_acc, val_loss = evaluate_one_epoch(epoch)
    return epoch, train_acc, train_loss, val_acc, val_loss

def main() -> None:
    """
    Main function to run the training and evaluation in parallel using multiprocessing.
    """
    lr = config["lr"]
    bs = config["bs"]
    epochs = config["epochs"]

    # Initialize Weights and Biases
    wandb.init(project="my_project", config=config)

    with ProcessPoolExecutor() as executor:
        futures = [executor.submit(run_epoch, epoch, lr, bs) for epoch in np.arange(1, epochs)]
        for future in as_completed(futures):
            epoch, train_acc, train_loss, val_acc, val_loss = future.result()
            wandb.log({
                "epoch": epoch,
                "train_acc": train_acc,
                "train_loss": train_loss,
                "val_acc": val_acc,
                "val_loss": val_loss
            })
            print(f"epoch: {epoch}")
            print(f"training accuracy: {train_acc}, training loss: {train_loss}")
            print(f"validation accuracy: {val_acc}, validation loss: {val_loss}")

if __name__ == "__main__":
    main()

问题解答 #

1. W&B多进程中count参数的作用及配置 #

count是W&B Sweep Agent的命令行参数，用来指定单个Agent进程要执行的实验次数，默认值为1。要实现最大化并行且无重复超参数：

• 用CLI启动Agent时，可通过--count N设置每个进程跑N个实验，比如wandb agent --count 5 sweep_id，让一个进程连续跑5个不同的超参数实验。
• 最大化并行的核心是启动与硬件资源匹配的Agent进程数（比如CPU核心数、GPU数量），每个进程由W&B调度器分配唯一的超参数组合——W&B调度器会自动跟踪已分配的超参数，确保不会重复分发，无需手动通过count控制去重。

2. 确保超参数唯一性 #

完全不需要自行在Python代码中处理，W&B Sweep内置了超参数分配与去重机制：

• Sweep的调度器（随机/网格/贝叶斯）会维护一个全局的超参数任务队列，每个Agent进程在调用wandb.init()时，会从队列中获取未被分配的超参数组合。
• 只要通过官方的Sweep流程创建任务、启动Agent，就能保证每个实验的超参数唯一，不会出现重复评估的情况。

3. 并行最佳实践 #

仅CPU多核场景

• 优先用CLI启动多Agent进程：直接通过bash脚本启动与CPU核心数一致的Agent进程，每个进程处理独立的超参数实验，避免Python多进程池的初始化开销。
• 若使用ProcessPoolExecutor，需确保子进程独立初始化W&B：不要在主进程中调用wandb.init()，而是在每个子进程的任务函数内调用wandb.init()获取Sweep的超参数，避免进程间的W&B状态冲突。
• 控制进程数不超过CPU核心数的80%，留出资源给系统进程，避免过度调度导致性能下降。

多GPU场景

• GPU与Agent进程一一绑定：为每个GPU启动一个独立的Agent进程，通过CUDA_VISIBLE_DEVICES环境变量指定GPU编号，确保每个实验独占一个GPU资源。
• 若实验本身需要多GPU分布式训练，可在单个Sweep实验内使用PyTorch Distributed等框架，同时让Sweep管理不同超参数实验的并行。
• 避免单个Agent进程占用多个GPU，除非是专门的分布式训练任务，否则会浪费GPU资源，且不利于多超参数实验的并行。

CLI/Bash实现方案 #

1. 定义Sweep配置文件（sweep_config.yaml） #

program: train.py
method: random  # 可选random/grid/bayes
parameters:
  lr:
    min: 1e-5
    max: 1e-3
  bs:
    values: [16, 32, 64]
  epochs:
    value: 5

2. 创建Sweep任务 #

wandb sweep sweep_config.yaml

执行后会得到一个Sweep ID，后续启动Agent需要用到。

3. CPU多核并行启动 #

# 启动4个Agent进程，每个进程跑3个实验
for i in {1..4}; do
  wandb agent --count 3 YOUR_SWEEP_ID &
done
wait

4. 多GPU并行启动 #

# 为GPU 0-3各启动一个Agent进程
for gpu in 0 1 2 3; do
  CUDA_VISIBLE_DEVICES=$gpu wandb agent YOUR_SWEEP_ID &
done
wait

适配Sweep的代码修改示例 #

把原脚本改成适配W&B Sweep的模式，让进程自动获取唯一超参数：

import random
import numpy as np
import wandb
from typing import Tuple, Dict

def train_one_epoch(epoch: int, lr: float, bs: int) -> Tuple[float, float]:
    """
    Simulate training for one epoch.
    """
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss

def evaluate_one_epoch(epoch: int) -> Tuple[float, float]:
    """
    Simulate evaluation for one epoch.
    """
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss

def run_epoch(epoch: int, lr: float, bs: int) -> Tuple[int, float, float, float, float]:
    """
    Run training and evaluation for one epoch.
    """
    train_acc, train_loss = train_one_epoch(epoch, lr, bs)
    val_acc, val_loss = evaluate_one_epoch(epoch)
    return epoch, train_acc, train_loss, val_acc, val_loss

def main() -> None:
    """
    Main function to run training/evaluation with W&B Sweep.
    """
    # 从W&B Sweep获取超参数配置
    run = wandb.init()
    config = run.config

    lr = config["lr"]
    bs = config["bs"]
    epochs = config["epochs"]

    for epoch in range(1, epochs + 1):
        epoch, train_acc, train_loss, val_acc, val_loss = run_epoch(epoch, lr, bs)
        wandb.log({
            "epoch": epoch,
            "train_acc": train_acc,
            "train_loss": train_loss,
            "val_acc": val_acc,
            "val_loss": val_loss
        })
        print(f"epoch: {epoch}")
        print(f"training accuracy: {train_acc}, training loss: {train_loss}")
        print(f"validation accuracy: {val_acc}, validation loss: {val_loss}")

if __name__ == "__main__":
    main()

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