如何用Python按Rank前缀合并同组.pt文件并生成唯一命名文件？

实现思路 #

1. 遍历目标目录，筛选出所有符合file_rankN_****.pt格式的文件
2. 按文件名中的rankN标识对文件进行分组
3. 加载每组内的所有.pt文件内容（基于PyTorch的torch.load），根据实际需求合并内容
4. 将合并后的内容保存为对应rank的file_rankN_combine.pt文件

完整代码实现 #

import os
import torch
from collections import defaultdict

# 目标目录（这里用当前目录，可根据实际修改）
target_dir = "./"

# 1. 按rank分组文件
rank_files = defaultdict(list)
for filename in os.listdir(target_dir):
    if filename.endswith(".pt") and "file_rank" in filename:
        # 提取rank编号，比如从"file_rank0_1011.pt"中得到"rank0"
        rank_part = filename.split("_")[1]
        rank_files[rank_part].append(os.path.join(target_dir, filename))

# 2. 遍历每个rank组，合并并保存文件
for rank, files in rank_files.items():
    combined_data = {}
    for file_path in files:
        # 加载单个.pt文件内容，指定cpu避免GPU内存占用问题
        data = torch.load(file_path, map_location="cpu")
        # 合并逻辑：默认处理PyTorch模型的state_dict字典（重复键会覆盖，按需调整）
        if isinstance(data, dict):
            combined_data.update(data)
        # 若存储的是张量，可替换为拼接逻辑：
        # if isinstance(data, torch.Tensor):
        #     combined_data.append(data)
        # 后续用torch.cat(combined_data)完成拼接

    # 生成并保存合并文件
    save_name = f"file_{rank}_combine.pt"
    save_path = os.path.join(target_dir, save_name)
    torch.save(combined_data, save_path)
    print(f"合并完成：{save_path}")

关键说明 #

• 分组逻辑：通过拆分文件名提取rankN标识，用字典自动完成文件分组
• 合并逻辑适配：默认处理PyTorch模型的state_dict字典，若你的.pt文件存储的是张量、自定义数据结构，需修改对应合并代码块
• 内存优化：加载文件时指定map_location="cpu"，避免大文件占用GPU内存
• 重复文件处理：同一rank下的重名文件（如示例中的file_rank2_1111.pt），合并时会覆盖之前的内容，可根据需求添加去重或累加逻辑

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