如何将PyTorch多维张量存储到磁盘并按偏移读取指定行数据

API 探路者

2026-10-6

最优实现方案

针对你这个固定形状的二维PyTorch张量，核心需求是随机读取指定行、无需加载全量数据，最优方案是采用「二进制裸数组存储+内存映射读取」的组合，比存成pt、pickle等需要全量加载的格式效率高很多，对更大尺寸的张量也通用。

第一步：存储张量到磁盘

先处理张量去掉梯度和cuda设备属性，分别存储张量元数据（形状、数据类型）和二进制数组：

import numpy as np
import torch
import json

# 替换为你自己的张量变量名
raw_tensor = tensor
# 转移到CPU、去掉梯度后转numpy数组
np_arr = raw_tensor.detach().cpu().numpy()

# 存储元数据，文件很小不占空间
metadata = {"shape": np_arr.shape, "dtype": str(np_arr.dtype)}
with open("emb_meta.json", "w", encoding="utf-8") as f:
    json.dump(metadata, f)

# 存储二进制张量文件
np_arr.tofile("entity_emb.bin")

第二步：读取指定行到内存

用numpy的memmap做内存映射，仅加载你需要的行数据，不会读取全量张量：

import numpy as np
import torch
import json

# 先读元数据
with open("emb_meta.json", "r", encoding="utf-8") as f:
    metadata = json.load(f)

# 只读模式映射文件，这一步不会加载全量数据到内存
mem_map = np.memmap(
    "entity_emb.bin",
    dtype=metadata["dtype"],
    shape=tuple(metadata["shape"]),
    mode="r"
)

# 直接索引你需要的行即可，比如最后一行，这时候才会把对应行加载到内存
target_row = mem_map[-1]
# 需要转回PyTorch张量的话直接转即可
target_tensor = torch.from_numpy(target_row.copy())

# 用完删除映射释放资源
del mem_map