PyTorch——Random相关状态管理

整体说明

  • PyTorch 中，包含很多随机操作，比如
    • 可以使用 torch.rand() 等函数获取随机数
    • 可以使用 torch.nn.functional.dropout() 实现随机 drop 一些神经元
    • 可以使用 tensor.random_() 等函数随机初始化参数
  • 这些涉及随机数/采样的方法均受限于一个随机状态管理

torch Seed 打印

  • torch Seed 打印代码： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    # 打印 torch 的随机种子情况
    def print_torch_seeds():
    	print("=" * 30 + "PyTorch Random Seeds Status")
    	print("=" * 30)
    	cpu_seed = torch.initial_seed()
    	print(f"[CPU] Seed: {cpu_seed}")

    	if torch.cuda.is_available():
    		try:
    			gpu_seed = torch.cuda.initial_seed()
    			current_device = torch.cuda.current_device()
    			device_name = torch.cuda.get_device_name(current_device)

    			print(f"[GPU] Seed: {gpu_seed}")
    			print(f"      Device: {current_device} ({device_name})")
    		except Exception as e:
    			print(f"[GPU] Error getting seed: {e}")
    	else:
    		print("[GPU] CUDA is not available.")

    	print("=" * 30)
    print_torch_seeds()

torch Seed 设置

  • 全局 torch Seed 设置代码：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    import torch

    # 固定CPU种子
    torch.manual_seed(42)

    # 固定所有GPU的种子（单GPU/多GPU通用）
    if torch.cuda.is_available():
    	torch.cuda.manual_seed_all(42)  # 替代 torch.cuda.manual_seed(42)（单GPU）

    # GPU上生成随机排列
    perm = torch.randperm(10, device="cuda") # 注意：需要指定 "cuda" 才会在 GPU 上执行
    print("GPU随机排列：", perm)  # 每次运行结果一致
    print("draw a random number：", torch.rand()) # 每次运行结果一致

  • 使用独立的 torch 生成器（独立管理自己的随机生成器）：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    import torch

    # 创建独立的生成器并设置种子
    generator = torch.Generator()
    generator.manual_seed(42)
    # generator = torch.Generator().manual_seed(42) # 等价实现

    # 生成随机排列时指定生成器
    perm1 = torch.randperm(10, generator=generator)
    perm2 = torch.randperm(10, generator=generator)

    print("独立生成器-第一次：", perm1)  # tensor([2, 7, 3, 1, 0, 9, 4, 5, 8, 6])
    print("独立生成器-第二次：", perm2)  # tensor([2, 0, 7, 9, 8, 4, 3, 6, 1, 5])

    # 重置生成器种子，结果重复
    generator.manual_seed(42)
    perm3 = torch.randperm(10, generator=generator)
    print("重置生成器后：", perm3)  # tensor([2, 7, 3, 1, 0, 9, 4, 5, 8, 6])（和perm1一致）
    • 说明：torch.Generator 是 PyTorch 中统一的随机数生成器（RNG）核心对象，几乎所有 PyTorch 内置的随机操作都支持通过 generator 参数指定该生成器

附录：torch.Generator 详细说明

  • torch.Generator 是 PyTorch 中统一的随机数生成器（RNG）核心对象 ，几乎所有 PyTorch 内置的随机操作都支持通过 generator 参数指定该生成器，仅极少数场景不支持（或无需支持）
  • torch.Generator 的核心作用是隔离随机状态 ：任何依赖 PyTorch 内置随机数生成的操作，只要支持 generator 参数，就能通过该生成器控制随机行为；无 generator 参数的操作，要么不依赖随机数，要么复用全局生成器（CPU/CUDA）
  • 所有需要随机逻辑的场景均支持 torch.Generator 的随机操作（全场景）
  • 注：无随机逻辑的操作本身无随机行为，因此不需要（也无法）指定 generator：
    • 张量基础操作：torch.ones()、torch.zeros()、torch.arange()、torch.cat()、torch.matmul() 等
    • 数学运算：torch.sin()、torch.exp()、torch.mean()、torch.argmax() 等
    • 索引/切片：x[:, 0]、x.index_select() 等
    • 设备/类型转换：x.to('cuda')、x.float() 等

随机逻辑的场景示例

  • 所有操作均可通过 generator 参数指定自定义 torch.Generator，实现随机状态隔离
  • 基础随机数生成
    函数/方法 用途 示例
    torch.rand() 均匀分布随机数 torch.rand(3, generator=g)
    torch.randn() 标准正态分布随机数 torch.randn(2, 4, generator=g)
    torch.randint() 整数随机数 torch.randint(0, 10, (3,), generator=g)
    torch.randperm() 随机排列 torch.randperm(5, generator=g)
    torch.rand_like() 按形状生成均匀随机数 torch.rand_like(torch.ones(2), generator=g)
    torch.randn_like() 按形状生成正态随机数 torch.randn_like(torch.ones(2), generator=g)
    torch.normal() 自定义均值/方差正态分布 torch.normal(0, 1, (3,), generator=g)
    torch.poisson() 泊松分布随机数 torch.poisson(torch.ones(3), generator=g)
    torch.exponential() 指数分布随机数 torch.exponential(1.0, (3,), generator=g)
    torch.cauchy() 柯西分布随机数 torch.cauchy(0, 1, (3,), generator=g)
    torch.log_normal() 对数正态分布随机数 torch.log_normal(0, 1, (3,), generator=g)
    torch.multinomial() 多项分布采样 torch.multinomial(torch.ones(5), 3, generator=g)
    torch.bernoulli() 伯努利分布（0/1） torch.bernoulli(torch.ones(3)*0.5, generator=g)
    • 注：指定参数 generator 时，前面的参数也需要指定（Python 本身的规则）
  • 张量随机初始化
    函数/方法 用途 示例
    tensor.random_() 原地随机初始化（整数） tensor.random_(generator=g)
    tensor.uniform_() 原地均匀分布初始化 tensor.uniform_(0, 1, generator=g)
    tensor.normal_() 原地正态分布初始化 tensor.normal_(0, 1, generator=g)
    tensor.cauchy_() 原地柯西分布初始化 tensor.cauchy_(0, 1, generator=g)
  • 随机采样/变换（数据增强等）
    函数/方法 用途 示例
    torch.utils.data.RandomSampler 数据集随机采样 RandomSampler(dataset, generator=g)
    torch.nn.functional.dropout() Dropout层随机失活 F.dropout(x, p=0.5, generator=g)
    torch.nn.functional.dropout2d() 2D Dropout F.dropout2d(x, p=0.5, generator=g)
    torch.nn.functional.dropout3d() 3D Dropout F.dropout3d(x, p=0.5, generator=g)
    torchvision.transforms 中的随机变换 图像随机增强（如RandomCrop） transforms.RandomCrop(32, generator=g)（需torchvision）
    torch.distributions 分布采样 概率分布采样（如Normal、Uniform） dist = Normal(0, 1); dist.sample((3,), generator=g)

特殊说明：随机场景但不支持 torch.Generator 的场景

  • 有随机逻辑但不支持自定义 generator 的场景；依赖随机数，但 PyTorch 未开放 generator 参数，只能复用全局生成器（CPU/CUDA）：
    操作 原因 替代方案
    torch.shuffle() 底层绑定全局生成器 用 torch.randperm(generator=g) 手动实现洗牌
    torch.nn.Dropout 模块（默认） 模块初始化时未绑定生成器 改用 F.dropout(generator=g) 或自定义模块绑定生成器
    部分第三方库的随机操作（如某些数据增强） 未适配 generator 参数 替换为 PyTorch 原生实现或手动设置全局种子
    torch.multiprocessing 多进程随机 进程间生成器隔离限制 每个进程内重新初始化 generator
  • 实践思考：
    • 1）凡是生成随机数的 PyTorch 原生函数，优先检查是否有 generator 参数，有则建议使用（隔离随机状态）
    • 2）对不支持 generator 的随机操作，要么手动实现（如用 randperm 替代 shuffle），要么临时设置全局种子并尽快恢复
    • 3）CUDA 场景务必创建对应设备的 generator，避免跨设备混用导致随机状态混乱

最佳实践：torch.Generator 隔离随机状态

  • 多个 torch.Generator 隔离随机状态示例 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    import torch

    # 创建两个独立生成器
    g1 = torch.Generator().manual_seed(42)
    g2 = torch.Generator().manual_seed(42)

    # 用g1生成随机数（消耗g1的状态）
    a = torch.rand(2, generator=g1)
    b = torch.rand(2, generator=g1)

    # 用g2生成随机数（g2状态未被消耗，结果和g1初始一致）
    c = torch.rand(2, generator=g2)

    print("a (g1第一次):", a)  # tensor([0.8823, 0.9150])
    print("b (g1第二次):", b)  # tensor([0.3829, 0.9593])
    print("c (g2第一次):", c)  # tensor([0.8823, 0.9150])（和a一致）

附录：GPU 下的 torch.Generator

  • torch.Generator 必须与操作的设备（CPU/CUDA）对齐 ，否则会导致隐式设备拷贝、性能损耗，甚至随机状态混乱
    • CPU 操作时使用 CPU 生成器
    • CUDA 操作时使用对应 CUDA 设备的生成器
    • 核心目的：避免隐式跨设备拷贝，保证随机状态的隔离性和可复现性
  • 所有支持 CUDA 的随机操作（如 torch.rand(3, device='cuda', generator=g)），需指定与生成器同设备的 generator
  • CUDA 生成器的随机状态与 CPU 生成器完全隔离，互不干扰
生成器的设备属性
  • torch.Generator 可通过 device 参数绑定具体设备，默认是 CPU： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    import torch

    # CPU生成器（默认）
    g_cpu = torch.Generator()  # 等价于 torch.Generator(device="cpu")
    print("CPU生成器设备：", g_cpu.device)  # 输出：cpu

    # CUDA生成器（需显式指定）
    if torch.cuda.is_available():
    	g_cuda = torch.Generator(device="cuda:0")  # 绑定cuda:0
    	print("CUDA生成器设备：", g_cuda.device)  # 输出：cuda:0
对齐 vs 不对齐的示例对比

  • 正确：生成器设备 等于 操作 device（推荐）

    1
    2
    3
    4
    5
    6
    7
    8
    if torch.cuda.is_available():
    	# 创建cuda:0的生成器
    	g_cuda = torch.Generator(device="cuda:0").manual_seed(42)
    	# 操作指定device=cuda:0，与生成器对齐
    	perm = torch.randperm(10, device="cuda:0", generator=g_cuda)
    	
    	print("结果设备：", perm.device)  # cuda:0
    	print("无隐式拷贝，效率最高")

  • 错误：生成器设备 不等于 操作 device（性能坑）

    1
    2
    3
    4
    5
    6
    7
    if torch.cuda.is_available():
    	# 生成器是cuda:0，但操作指定 device=cpu
    	g_cuda = torch.Generator(device="cuda:0").manual_seed(42)
    	perm = torch.randperm(10, device="cpu", generator=g_cuda)
    	
    	print("结果设备：", perm.device)  # cpu
    	print("隐式拷贝：GPU生成随机数 拷贝到CPU（额外开销）")

  • 更隐蔽的错误：CUDA操作用CPU生成器

    1
    2
    3
    4
    5
    6
    7
    if torch.cuda.is_available():
    	# 生成器是CPU，操作指定device=cuda
    	g_cpu = torch.Generator().manual_seed(42)
    	perm = torch.randperm(10, device="cuda", generator=g_cpu)
    	
    	print("结果设备：", perm.device)  # cuda:0
    	print("隐式拷贝：CPU生成随机数 拷贝到GPU（额外开销）")
为什么必须对齐？
  • 随机数生成器的硬件绑定 ：
    • CPU 生成器依赖 CPU 的随机数算法
    • CUDA 生成器依赖 GPU 的 cuRAND 库，直接在 GPU 显存生成随机数；
      • 跨设备使用时，PyTorch 会先在生成器设备生成随机数，再通过 PCIe 总线拷贝到操作指定的设备，产生额外耗时
  • 随机状态的隔离性（容易因为误用而出错） ：
    • CUDA生成器的随机状态（get_state()）和 CPU 生成器完全隔离，若跨设备使用，会导致“生成器状态和操作设备不匹配”，破坏随机种子的可复现性：
      1
      2
      3
      4
      5
      6
      7
      if torch.cuda.is_available():
      	g_cuda = torch.Generator(device="cuda").manual_seed(42)
      	# 第一次：跨设备使用（cuda生成器 到 cpu操作）
      	perm1 = torch.randperm(10, device="cpu", generator=g_cuda)
      	# 第二次：直接用cuda生成器 到 cuda操作
      	perm2 = torch.randperm(10, device="cuda", generator=g_cuda)
      	# perm2的结果不等于“重新seed后cuda操作的结果”（状态已被跨设备操作消耗）
torch.Generator 的最佳实践

  • 创建生成器时显式指定设备 ：

    • 不要依赖默认的CPU生成器，GPU场景务必创建 device="cuda" 的生成器

  • 封装成函数，强制对齐 ：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    def get_generator(device: str = "cpu", seed: int = 42):
        g = torch.Generator(device=device)
        g.manual_seed(seed)
        return g

    # 使用
    if torch.cuda.is_available():
        g = get_generator(device="cuda:0", seed=42)
        perm = torch.randperm(10, device="cuda:0", generator=g)

  • 多GPU场景：每个GPU对应独立生成器 ：

    1
    2
    3
    4
    5
    6
    7
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        # 为cuda:0和cuda:1分别创建生成器
        g0 = torch.Generator(device="cuda:0").manual_seed(42)
        g1 = torch.Generator(device="cuda:1").manual_seed(100)
        
        perm0 = torch.randperm(10, device="cuda:0", generator=g0)
        perm1 = torch.randperm(10, device="cuda:1", generator=g1)