如何在PyTorch中高效存储种子向量以保证计算可复现？

问题核心 #

你当前的方案存在两个关键问题，导致循环中生成的向量重复：

1. 浮点数种子的精度丢失：torch.manual_seed()仅接受整数类型参数，传入浮点数会被强制截断为整数，不同浮点数可能对应同一个整数种子，最终生成相同的随机向量。
2. 循环种子逻辑错误：若循环内重复使用同一种子，或种子生成方式未保证唯一性，会导致每次重置随机数生成器（RNG）到相同状态，进而生成重复向量。

你的需求是合理的：用单个8字节种子代替大体积随机向量（如64KB）来复现单张图像生成，大幅节省存储。以下是修正后的实现方案。

正确实现步骤 #

1. 使用整数种子代替浮点数 #

整数种子无精度丢失问题，通过torch.randint()可生成范围足够大的唯一整数种子，确保每张图像对应独立的RNG状态。

2. 循环生成并存储独立种子 #

为每张图像生成唯一整数种子并存储；生成向量时，用对应种子初始化RNG，确保向量可复现。

代码示例：生成阶段 #

import torch

num_images = 10
vector_shape = (8192,)  # 对应你提到的64KB向量（8192*8字节=65536字节）

# 可选：设置全局初始种子，确保整个种子序列也可复现（不需要可复现种子序列可跳过）
# torch.manual_seed(42)
# 如果用GPU，需同步设置CUDA种子
# torch.cuda.manual_seed_all(42)

for img_idx in range(num_images):
    # 生成唯一的int64类型种子，覆盖所有整数取值范围
    seed = torch.randint(0, 2**63 - 1, (1,), dtype=torch.int64).item()
    # 用种子初始化RNG
    torch.manual_seed(seed)
    # 如果用GPU，同步设置CUDA种子
    # torch.cuda.manual_seed(seed)
    
    # 生成对应随机向量
    myvector = torch.randn(vector_shape)
    # 存储种子（可选择更高效的批量存储方式，见下文）
    torch.save(seed, f"seed_img_{img_idx}.pt")
    
    # 此处添加生成图像的逻辑...

代码示例：复现阶段 #

import torch

vector_shape = (8192,)
# 加载目标图像对应的种子
target_seed = torch.load("seed_img_3.pt")
# 初始化RNG
torch.manual_seed(target_seed)
# 如果用GPU，同步设置
# torch.cuda.manual_seed(target_seed)

# 生成与原过程完全相同的随机向量
myvector = torch.randn(vector_shape)
# 重新生成图像...

优化存储效率 #

若生成大量图像，单独存储每个种子文件会有额外开销，可将所有种子存入数组后批量保存：

# 生成阶段批量存储种子
seeds = []
for img_idx in range(num_images):
    seed = torch.randint(0, 2**63 - 1, (1,), dtype=torch.int64).item()
    seeds.append(seed)
    # ...生成向量和图像的逻辑...

# 批量保存所有种子，仅需8*num_images字节的实际数据
torch.save(torch.tensor(seeds), "all_seeds.pt")

# 复现时加载批量种子
all_seeds = torch.load("all_seeds.pt")
target_seed = all_seeds[3].item()
torch.manual_seed(target_seed)
myvector = torch.randn(vector_shape)

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