一、问题背景:为什么大对象序列化会让内存爆掉

当你需要用 numpy 把一个大数组存到磁盘或者从磁盘加载回来的时候,可能会遇到一个很头疼的问题:内存占用突然翻倍甚至翻三倍。明明你的数据文件只有几百兆,加载到内存里却占了好几个 GB,电脑直接卡死,进程被系统杀掉。

这种情况在训练深度学习模型参数、处理大规模图像数据集、存储数值模拟的计算结果时非常常见。np.save 和 np.load 是 numpy 提供的最基础的数组序列化接口,用起来简单,但在处理大对象时,如果不了解它的底层机制,很容易踩坑。

这篇文章就来聊聊怎么解决这些问题,以及有哪些替代方案可以在不牺牲太多性能的前提下大幅降低内存占用。

二、np.save 和 np.load 的底层机制

2.1 np.save 是怎么工作的

np.save 的原理其实很直接:它把 numpy 数组的头信息(包括数据类型、形状、字节序等)和原始的内存数据一起写入到一个二进制文件里。这个文件就是 .npy 格式的文件。

# 技术栈:Python + NumPy

import numpy as np

# 创建一个 10000x10000 的浮点数组,大约占 800MB
arr = np.random.rand(10000, 10000)

# 保存到文件,整个数组会被一次性写入磁盘
# 内部会把内存数据完整拷贝一次到文件缓冲区
np.save("large_array.npy", arr)

这里面的关键点是:np.save 在执行时会把内存中的数组数据完整地拷贝一次到磁盘文件中。这个拷贝过程本身就需要额外的内存空间来充当缓冲区。

2.2 np.load 是怎么工作的

np.load 则是反过来,把 .npy 文件从磁盘读回到内存中。默认情况下,它会读取整个文件内容,然后在内存中创建一个完整的数组副本。

# 从文件加载数组
# 默认行为:读取整个文件到内存中,创建一个完整的数组
loaded_arr = np.load("large_array.npy")

# 此时 loaded_arr 和磁盘上的文件内容完全一致
# 但内存中已经有了一个完整的副本

问题就出在这里:如果你在加载的时候内存里已经有一份原始数据,那么加载操作会再创建一份完整的副本,内存占用直接翻倍。

三、内存暴涨的核心原因分析

3.1 双重缓冲问题

np.save 在写入文件时,numpy 内部会使用缓冲区来暂存数据。对于大数组来说,这个缓冲区可能会非常占用内存。在某些系统配置下,Python 的垃圾回收机制可能不会及时回收旧对象,导致新旧数据同时存在于内存中。

3.2 类型转换的隐式开销

当你加载一个 .npy 文件时,numpy 可能会进行类型转换。比如文件中的数据是 float32,但你的代码期望的是 float64,numpy 会创建一个新的数组来存储转换后的结果,旧数组在转换完成之前不会被释放。

3.3 Python 对象头的额外开销

除了数据本身,Python 的数组对象在内存中还会有对象头(object header)的开销。对于多维数组来说,每个维度都可能有额外的元数据开销。当数组非常大时,这些额外开销加起来也不容忽视。

# 技术栈:Python + NumPy
# 演示内存暴涨的具体表现

import numpy as np
import tracemalloc

# 开启内存追踪
tracemalloc.start()

# 创建一个约 800MB 的数组
big_array = np.random.rand(10000, 10000)

snapshot1 = tracemalloc.take_snapshot()

# 保存操作:内部会创建缓冲区,内存可能暂时翻倍
np.save("test_large.npy", big_array)

snapshot2 = tracemalloc.take_snapshot()

# 比较两个快照的内存使用情况
top_stats = snapshot2.compare_to(snapshot1, "lineno")
for stat in top_stats[:5]:
    print(f"内存增长:{stat.count} 块对象,占用 {stat.size / 1024 / 1024:.2f} MB")

上面的代码展示了用 tracemalloc 追踪内存变化的方法。可以看到,在 np.save 执行过程中,内存占用会出现明显的峰值。

四、优化方案一:使用内存映射(mmap_mode)

4.1 什么是内存映射

内存映射是一种让操作系统直接管理文件到内存映射的技术。简单来说,它不是把文件完整地读到内存里,而是告诉操作系统:"这个文件的某个区域,你帮我映射到内存地址空间里,我需要什么你就给我什么。"

这样做的好处是:操作系统会根据实际的内存使用情况,只把需要的数据加载到物理内存中,而不是提前加载全部数据。

4.2 在 np.load 中使用 mmap_mode

# 技术栈:Python + NumPy
# 使用内存映射加载大数组,避免一次性加载全部数据到内存

import numpy as np

# 先创建一个测试用大数组并保存
large_data = np.random.rand(50000, 50000)
np.save("huge_array.npy", large_data)

# 方式一:只读模式加载(默认只读,不会把数据写入原文件)
# mmap_mode='r' 表示以只读方式映射到内存
mapped_array = np.load("huge_array.npy", mmap_mode='r')

# 此时并没有把整个文件加载到内存中!
# 只有你实际访问的元素才会被加载
print(f"数组形状:{mapped_array.shape}")

# 只读取一个小的切片,只有这部分数据会被加载到内存
subset = mapped_array[0:100, 0:100]  # 只加载 100x100 的数据
print(f"切片数据的前5个值:{subset[:5]}")

# 方式二:如果需要对加载的数据进行修改并写回原文件
# mmap_mode='r+' 表示以读写模式映射
mapped_array_rw = np.load("huge_array.npy", mmap_mode='r+')
mapped_array_rw[0, 0] = 999.0  # 修改会被写回原文件

# 方式三:mmap_mode='c' 复制模式
# 会创建文件的副本,对副本的修改不影响原文件
mapped_array_copy = np.load("huge_array.npy", mmap_mode='c')

4.3 mmap_mode 的使用场景

mmap_mode 特别适合以下场景:你需要反复访问同一个大文件的不同部分、你只需要文件中的某一部分数据、多个进程需要共享同一个大文件的数据。

# 技术栈:Python + NumPy
# mmap_mode 在迭代处理大文件时的应用

import numpy as np

# 假设我们有一个非常大的训练数据集,按行存储
# 每行是一个样本,共有 1000 万个样本
dataset_path = "training_data.npy"

# 使用 mmap 方式加载
dataset = np.load(dataset_path, mmap_mode='r')

# 分批处理,每批 1000 个样本
batch_size = 1000
num_batches = len(dataset) // batch_size

for i in range(num_batches):
    # 每次只读取一个 batch 的数据
    # 操作系统会自动管理内存,用完的 batch 可以被释放
    batch = dataset[i * batch_size: (i + 1) * batch_size]
    
    # 对 batch 进行处理(比如送入模型训练)
    loss = compute_loss(batch)
    print(f"批次 {i + 1}/{num_batches},损失:{loss:.4f}")

五、优化方案二:分块读写策略

5.1 为什么要分块

即使使用了 mmap_mode,有些场景下你仍然需要把整个文件加载到内存中进行计算(比如需要全局的统计信息、需要做全局排序等)。这时候,分块处理就是一个有效的策略。

分块的核心思想是:把一个大数组拆分成多个小块,每块都在内存能承受的范围内,然后逐块处理。

5.2 实现分块读写

# 技术栈:Python + NumPy
# 实现大数组的分块读写工具

import numpy as np
import os

def chunked_save(array, filepath, chunk_size=None):
    """
    分块保存大数组到文件
    
    参数:
        array: numpy 数组
        filepath: 保存的文件路径
        chunk_size: 每次处理的行数,None 表示自动计算
    """
    if chunk_size is None:
        # 根据数组大小自动计算合适的块大小
        # 目标:每次处理的块不超过 256MB
        target_bytes = 256 * 1024 * 1024
        element_size = array.itemsize
        rows_in_chunk = target_bytes // (element_size * array.shape[1])
        chunk_size = max(1, rows_in_chunk)
    
    num_chunks = array.shape[0] // chunk_size
    
    with open(filepath, 'wb') as f:
        for i in range(num_chunks):
            start_idx = i * chunk_size
            end_idx = (i + 1) * chunk_size
            
            # 逐块写入
            chunk = array[start_idx:end_idx]
            
            # 每块数据独立保存
            chunk_bytes = np.array(chunk).tobytes()
            f.write(chunk_bytes)
            
            print(f"已写入第 {i + 1}/{num_chunks} 块,"
                  f"当前内存使用:{chunk.nbytes / 1024 / 1024:.1f} MB")

def chunked_load(filepath, original_shape, dtype=np.float64, chunk_size=1000):
    """
    分块加载大数组
    
    参数:
        filepath: 文件路径
        original_shape: 原始数组的形状
        dtype: 数据类型
        chunk_size: 每次读取的行数
    """
    total_elements = np.prod(original_shape)
    
    result = np.empty(original_shape, dtype=dtype)
    
    file_element_size = np.dtype(dtype).itemsize
    elements_per_chunk = chunk_size * original_shape[1]
    
    with open(filepath, 'rb') as f:
        chunk_index = 0
        while True:
            chunk_data = f.read(elements_per_chunk * file_element_size)
            if not chunk_data:
                break
            
            chunk = np.frombuffer(chunk_data, dtype=dtype)
            chunk = chunk.reshape(chunk_size, original_shape[1])
            
            # 将块数据填入结果数组的对应位置
            result[chunk_index * chunk_size:
                   (chunk_index + 1) * chunk_size, :] = chunk
            
            chunk_index += 1
    
    return result

# 使用示例
if __name__ == "__main__":
    # 创建一个测试大数组
    data = np.random.rand(100000, 5000)  # 约 4GB
    
    # 分块保存
    chunked_save(data, "large_data.bin", chunk_size=1000)
    
    # 分块加载
    loaded_data = chunked_load("large_data.bin", 
                               original_shape=(100000, 5000),
                               dtype=np.float64,
                               chunk_size=1000)
    
    # 验证数据一致性
    print(f"数据是否一致:{np.allclose(data, loaded_data)}")

5.3 使用 numpy 的内存视图优化分块

# 技术栈:Python + NumPy
# 利用 numpy 的内存视图(view)来减少分块拷贝

import numpy as np

def view_based_chunking(arr, chunk_size):
    """
    利用 numpy 的 view 机制进行零拷贝分块
    注意:view 不会创建新的数据副本,只是重新解释内存
    """
    n_rows = arr.shape[0]
    n_chunks = (n_rows + chunk_size - 1) // chunk_size
    
    for i in range(n_chunks):
        start = i * chunk_size
        end = min(start + chunk_size, n_rows)
        
        # 创建视图,不拷贝数据
        # 这只是对原始数据的一种"解读方式"
        chunk_view = arr[start:end, :]
        
        # 对视图进行操作
        # 比如:计算每一块的统计信息
        block_mean = chunk_view.mean()
        block_std = chunk_view.std()
        
        print(f"块 {i}(行 {start}-{end}):"
              f"均值 = {block_mean:.4f},标准差 = {block_std:.4f}")

# 测试
test_data = np.random.rand(10000, 1000)
view_based_chunking(test_data, chunk_size=1000)

六、优化方案三:压缩存储与惰性加载

6.1 np.savez 和 np.savez_compressed

numpy 还提供了 np.savez 和 np.savez_compressed 来保存多个数组。np.savez_compressed 会对每个数组进行压缩后再保存,可以显著减小文件体积。

# 技术栈:Python + NumPy
# 使用压缩存储减小文件大小

import numpy as np
import os

# 创建一个稀疏的数据矩阵(大量零值,压缩率会很高)
sparse_matrix = np.zeros((10000, 10000))
sparse_matrix[0:1000, 0:1000] = np.random.rand(1000, 1000)

# 不压缩保存
np.savez("sparse_no_compress.npz", data=sparse_matrix)
size_no_compress = os.path.getsize("sparse_no_compress.npz")

# 压缩保存
np.savez_compressed("sparse_compressed.npz", data=sparse_matrix)
size_compressed = os.path.getsize("sparse_compressed.npz")

print(f"未压缩文件大小:{size_no_compress / 1024 / 1024:.2f} MB")
print(f"压缩后文件大小:{size_compressed / 1024 / 1024:.2f} MB")
print(f"压缩率:{size_compressed / size_no_compress * 100:.1f}%")

# 加载压缩的文件
# 使用 mmap_mode 可以减少内存占用
compressed_data = np.load("sparse_compressed.npz", mmap_mode='r')["data"]
print(f"加载后的数据形状:{compressed_data.shape}")

6.2 自定义惰性加载器

# 技术栈:Python + NumPy
# 自定义惰性加载器:只在需要时才加载数据

import numpy as np

class LazyArrayLoader:
    """
    惰性加载器:只在访问数据时才真正加载到内存
    不访问的数据不会占用内存
    """
    def __init__(self, filepath):
        self.filepath = filepath
        self._array = None
        self._file_obj = None
    
    @property
    def shape(self):
        """获取形状而不加载数据(通过 mmap 模式读取头信息)"""
        if self._array is None:
            self._file_obj = np.load(self.filepath, mmap_mode='r')
            return self._file_obj.shape
        return self._array.shape
    
    @property
    def dtype(self):
        if self._array is None:
            if self._file_obj is None:
                self._file_obj = np.load(self.filepath, mmap_mode='r')
            return self._file_obj.dtype
        return self._array.dtype
    
    def load(self):
        """真正加载数据到内存"""
        if self._array is None:
            self._array = np.load(self.filepath)
        return self._array
    
    def get_slice(self, indices):
        """获取数据的一个切片(通过 mmap 方式,不加载全部数据)"""
        if self._file_obj is None:
            self._file_obj = np.load(self.filepath, mmap_mode='r')
        return self._file_obj[indices]
    
    def __del__(self):
        """清理资源"""
        self._array = None
        self._file_obj = None

# 使用示例
if __name__ == "__main__":
    big_data = np.random.rand(10000, 10000)
    np.save("lazy_data.npy", big_data)
    
    # 使用惰性加载器
    lazy = LazyArrayLoader("lazy_data.npy")
    
    # 获取元数据(不加载数据)
    print(f"形状:{lazy.shape}")
    print(f"数据类型:{lazy.dtype}")
    
    # 获取一个小的切片
    small_slice = lazy.get_slice([0:10, 0:10])
    print(f"切片:{small_slice}")
    
    # 当确实需要全部数据时再加载
    full_data = lazy.load()

七、优化方案四:使用 HDF5 替代 npy 格式

7.1 为什么选择 HDF5

HDF5 是一种用于存储大量科学数据的文件格式。相比于 .npy 格式,HDF5 支持分块存储、压缩、并行读写等高级功能,在处理大规模数据时更加灵活和高效。

使用 h5py 库,你可以像操作 numpy 数组一样操作 HDF5 文件中的数据,同时还支持内存映射。

7.2 用 h5py 实现分块存储

# 技术栈:Python + NumPy + h5py
# 使用 HDF5 实现高效的大对象存储

import numpy as np
import h5py

def save_to_h5_chunks(arr, filepath, chunk_shape=None):
    """
    将大数组分块保存到 HDF5 文件
    
    参数:
        arr: numpy 数组
        filepath: HDF5 文件路径
        chunk_shape: 每个块的形状,None 表示自动计算
    """
    if chunk_shape is None:
        # 自动计算块大小:目标每个块约 1MB
        target_chunk_bytes = 1024 * 1024
        element_bytes = arr.itemsize
        elements_per_chunk = target_chunk_bytes // element_bytes
        
        if len(arr.shape) == 2:
            rows = max(1, elements_per_chunk // arr.shape[1])
            chunk_shape = (rows, arr.shape[1])
        else:
            chunk_shape = arr.shape
    
    print(f"数据形状:{arr.shape},块大小:{chunk_shape}")
    
    with h5py.File(filepath, 'w') as h5f:
        # 创建数据集,指定块大小
        # 注意:chunks 参数启用分块存储
        dset = h5f.create_dataset(
            "data",
            shape=arr.shape,
            dtype=arr.dtype,
            chunks=chunk_shape,
            compression="gzip",
            compression_opts=4
        )
        
        # 分块写入数据
        for i in range(0, arr.shape[0], chunk_shape[0]):
            end = min(i + chunk_shape[0], arr.shape[0])
            dset[i:end, :] = arr[i:end, :]
            print(f"  已写入 {end}/{arr.shape[0]} 行")

def load_from_h5_with_mmap(filepath, dataset_name="data"):
    """
    使用内存映射方式加载 HDF5 数据
    只有实际访问的数据才会被加载到内存
    """
    with h5py.File(filepath, 'r') as h5f:
        dataset = h5f[dataset_name]
        return dataset

# 使用示例
if __name__ == "__main__":
    data = np.random.rand(100000, 10000)
    
    # 分块保存到 HDF5
    save_to_h5_chunks(data, "large_data.h5", chunk_shape=(1000, 10000))
    
    # 使用内存映射加载(不一次性加载全部数据)
    dset = load_from_h5_with_mmap("large_data.h5")
    print(f"HDF5 数据集形状:{dset.shape}")
    
    # 只读取需要的部分
    subset = dset[0:100, :]
    print(f"子集数据前 3 行:{subset[:3]}")
    
    del dset

7.3 HDF5 与 npy 的性能对比

# 技术栈:Python + NumPy + h5py
# 对比 npy 和 HDF5 的读写性能

import numpy as np
import h5py
import time
import os

def bench_npy(filepath, data):
    """测试 npy 格式的读写性能"""
    start = time.time()
    np.save(filepath, data)
    write_time = time.time() - start
    
    start = time.time()
    loaded = np.load(filepath)
    read_time = time.time() - start
    
    file_size = os.path.getsize(filepath)
    del loaded
    
    return {
        "写入时间": f"{write_time:.3f} 秒",
        "读取时间": f"{read_time:.3f} 秒",
        "文件大小": f"{file_size / 1024 / 1024:.2f} MB"
    }

def bench_h5(filepath, data, chunk_size=1000):
    """测试 HDF5 格式的读写性能"""
    start = time.time()
    with h5py.File(filepath, 'w') as h5f:
        dset = h5f.create_dataset(
            "data",
            shape=data.shape,
            dtype=data.dtype,
            chunks=(chunk_size, data.shape[1]),
            compression="gzip"
        )
        dset[...] = data
    write_time = time.time() - start
    
    start = time.time()
    with h5py.File(filepath, 'r') as h5f:
        loaded = h5f["data"][:]
    read_time = time.time() - start
    
    file_size = os.path.getsize(filepath)
    del loaded
    
    return {
        "写入时间": f"{write_time:.3f} 秒",
        "读取时间": f"{read_time:.3f} 秒",
        "文件大小": f"{file_size / 1024 / 1024:.2f} MB"
    }

# 运行基准测试
if __name__ == "__main__":
    data = np.random.rand(10000, 5000)
    
    print("=== npy 格式 ===")
    npy_result = bench_npy("test.npy", data)
    for k, v in npy_result.items():
        print(f"  {k}: {v}")
    
    print("\n=== HDF5 格式 ===")
    h5_result = bench_h5("test.h5", data)
    for k, v in h5_result.items():
        print(f"  {k}: {v}")
    
    os.remove("test.npy")
    os.remove("test.h5")

八、优化方案五:零拷贝与共享内存

8.1 利用 numpy 的 frombuffer 和 tobytes

# 技术栈:Python + NumPy
# 使用 frombuffer 和 tobytes 实现零拷贝数据传递

import numpy as np

def zero_copy_pipeline(data):
    """
    演示零拷贝数据处理流程
    """
    # 将数组转换为字节缓冲区
    buffer = data.tobytes()
    print(f"缓冲区大小:{len(buffer) / 1024 / 1024:.2f} MB")
    
    # 从缓冲区创建新的数组视图
    # 注意:frombuffer 创建的是对同一内存的引用
    reconstructed = np.frombuffer(buffer, dtype=data.dtype)
    reconstructed = reconstructed.reshape(data.shape)
    
    print(f"原始数据前5个值:{data[:5]}")
    print(f"重建数据前5个值:{reconstructed[:5]}")
    print(f"数据一致:{np.allclose(data, reconstructed)}")

# 使用 np.memmap 实现真正的文件级零拷贝
def file_zero_copy(filepath, dtype=np.float64, shape=None):
    """
    使用 np.memmap 实现文件级零拷贝
    数据不会一次性加载到内存中
    """
    if shape is not None:
        # 创建新的内存映射文件
        mmapped = np.memmap(filepath, dtype=dtype, mode='w+', shape=shape)
    else:
        # 打开已有的内存映射文件
        mmapped = np.memmap(filepath, dtype=dtype, mode='r')
    
    return mmapped

# 使用示例
if __name__ == "__main__":
    # 创建一个内存映射文件
    mmapped = np.memmap(
        "data_memmap.dat",
        dtype='float64',
        mode='w+',
        shape=(10000, 1000)  # 约 80MB 的数组
    )
    
    # 直接操作内存映射数组(数据按需加载到内存)
    mmapped[0:100, :] = np.random.rand(100, 1000)
    mmapped.flush()  # 将数据刷写到磁盘
    
    # 重新打开文件(数据不会全部加载)
    mmapped2 = np.memmap("data_memmap.dat", dtype='float64', mode='r')
    print(f"内存映射数组形状:{mmapped2.shape}")
    print(f"前 100 行的均值:{mmapped2[0:100, :].mean():.4f}")
    
    del mmapped
    del mmapped2

九、不同方案的应用场景对比

9.1 方案适用场景总结

# 技术栈:Python
# 根据场景选择合适的方案

def recommend_approach(scenario, data_size_mb, access_pattern="sequential"):
    """
    根据场景推荐最优方案
    
    参数:
        scenario: 场景描述
        data_size_mb: 数据大小(MB)
        access_pattern: 访问模式 - sequential/random/partial/global
    """
    recommendations = []
    
    if data_size_mb > 5000:
        if access_pattern == "partial":
            recommendations.append("使用 np.load(mmap_mode='r') + 切片访问")
            recommendations.append("使用 np.memmap 进行文件级零拷贝")
        elif access_pattern == "sequential":
            recommendations.append("使用分块处理(chunk_size 根据可用内存调整)")
        elif access_pattern == "global":
            recommendations.append("使用 HDF5 + 分块读写")
            recommendations.append("考虑使用 out-of-core 计算框架(如 dask)")
    
    if data_size_mb <= 5000:
        if access_pattern == "partial":
            recommendations.append("mmap_mode='r' 足够")
        else:
            recommendations.append("标准 np.load/np.save 即可,但注意避免双重加载")
    
    return recommendations

# 使用示例
print("场景1 - 加载训练数据(8GB,部分访问):")
for r in recommend_approach("加载训练数据", 8000, "partial"):
    print(f"  {r}")

print("\n场景2 - 存储模型参数(200MB,全局访问):")
for r in recommend_approach("存储模型参数", 200, "global"):
    print(f"  {r}")

print("\n场景3 - 处理图像批量(5GB,顺序访问):")
for r in recommend_approach("处理图像批量", 5000, "sequential"):
    print(f"  {r}")

十、技术优缺点与注意事项

10.1 各方案的优缺点

mmap_mode 方案:

  • 优点:内存占用极低,支持随机访问,实现简单
  • 缺点:需要文件系统的支持,写入性能可能不如直接写入
  • 注意事项:在 Windows 系统上,mmap_mode 的行为可能与 Linux 不同,建议测试确认

分块处理方案:

  • 优点:内存使用完全可控,适用于任意大小的数据
  • 缺点:实现复杂度较高,需要处理块边界问题
  • 注意事项:分块大小时要平衡内存占用和读取效率,太小的块会导致频繁的磁盘 I/O

压缩存储方案:

  • 优点:显著减小文件体积,节省磁盘空间
  • 缺点:解压需要额外的 CPU 时间
  • 注意事项:压缩和解压的时间开销需要考虑,对于频繁读写的小文件,压缩可能不划算

HDF5 方案:

  • 优点:功能丰富,支持并行读写,支持多种压缩算法
  • 缺点:需要安装 h5py 库,文件结构比 .npy 复杂
  • 注意事项:HDF5 的锁机制可能会影响多进程并发访问,需要合理设置访问权限

10.2 常见陷阱和注意事项

  1. 内存泄漏:使用 mmap 或 memmap 时,一定要在不再使用时显式释放引用(del 变量),否则内存映射区域不会自动释放。

  2. 数据类型匹配:从文件加载数据时,确保预期的 dtype 与文件中的数据类型一致,否则可能导致隐式类型转换和额外的内存开销。

  3. 字节序问题:跨平台传输 .npy 文件时,注意字节序(endian)的差异。numpy 的 .npy 格式头中包含了字节序信息,通常会自动处理。

  4. 文件损坏:在写入大文件时如果中途断电或进程崩溃,可能会导致文件损坏。建议先写入临时文件,写入完成后原子性地重命名。

# 技术栈:Python + NumPy
# 安全的文件写入:使用临时文件 + 原子重命名

import numpy as np
import os
import tempfile

def safe_save(array, filepath):
    """
    安全地保存数组到文件
    使用临时文件 + 原子重命名来避免写入中途出错导致文件损坏
    """
    # 创建临时文件(与目标文件在同一目录下,确保在同一文件系统上)
    dir_name = os.path.dirname(filepath) or '.'
    tmp_fd, tmp_path = tempfile.mkstemp(suffix='.npy', dir=dir_name)
    
    try:
        os.close(tmp_fd)
        np.save(tmp_path, array)
        
        # 确保数据刷写到磁盘
        with open(tmp_path, 'rb') as f:
            os.fsync(f.fileno())
        
        # 原子性地重命名(在同一文件系统上是原子操作)
        os.replace(tmp_path, filepath)
        print(f"成功保存到 {filepath}")
    except Exception as e:
        # 如果出错,清理临时文件
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        raise e

# 使用示例
data = np.random.rand(10000, 10000)
safe_save(data, "safe_array.npy")
  1. numpy 版本兼容:np.load 和 np.save 的 .npy 文件格式在不同 numpy 版本之间可能不兼容。建议固定 numpy 版本,或者在保存时记录版本信息。

  2. 线程安全:np.load 和 np.save 本身是线程安全的,但如果你同时修改同一个数组对象,需要注意加锁或使用其他同步机制。

  3. 内存上限估算:在开始处理大文件之前,先用 arr.nbytes / (1024**2) 估算数组的内存占用,确认机器是否有足够的可用内存。如果数组大小超过可用内存的 80%,就应该考虑使用 mmap 或分块方案了。

十一、总结

处理大规模 numpy 数组的序列化问题,核心思路是"不要一次性把所有数据都塞进内存"。下面是一个快速参考的优化清单:

  • 如果只是读取文件的一部分数据,用 mmap_mode='r' 加载,然后通过切片按需访问,内存占用可以控制在极低的水平。
  • 如果需要对全部数据进行计算,使用分块处理策略,自己控制每次处理的数据量,确保不超过可用内存。
  • 如果数据有压缩需求,使用 np.savez_compressed 或者 HDF5 的压缩功能,在磁盘空间和读取速度之间找到平衡。
  • 如果需要更高级的功能,比如多个数据集、元数据属性、并行读写,HDF5 是比 .npy 更好的选择。
  • 如果需要多个进程共享数据,np.memmap 是一个非常好的选择,它利用了操作系统的虚拟内存机制来共享数据。

最后,无论选择哪种方案,都要记住几个关键点:及时释放不再使用的内存引用(使用 del)、在安全写入时使用临时文件加原子重命名、在跨平台场景下注意字节序问题。这些看似小的细节,往往决定了你的程序在真实生产环境中能否稳定运行。