Google Colab循环读取.h5文件时内存占用过高且程序冻结

问题描述 #

• 本地运行脚本正常，但在Google Colab中循环处理多个.h5文件时出现异常：
  • 前1-4个文件可成功处理
  • 内存占用突然飙升后回落
  • 峰值后程序在某一迭代处无限卡顿，无明确报错仅进程挂起
  • 文件顺序不影响问题发生，异常出现随机

已尝试操作 #

• 单独读取所有文件，验证文件本身均有效
• 调整文件处理顺序，问题仍随机出现
• 监控Colab内存，确认冻结前存在大幅内存飙升
• 尝试手动清理变量（del variable + gc.collect()），以及使用tables.open_file(..., driver="core")，均无改善

疑问 #

• Colab是否存在内存泄漏或HDF5缓存问题？
• h5py/tables在文件关闭后是否仍占用内存？
• 能否强制Colab在每次迭代后释放内存？
• 有无类似问题的解决方案？

代码片段 #

def read_h5(file_name):
    """ Reads an HDF5 file containing a sparse matrix in CSR format. """

    with tables.open_file(file_name, 'r') as f:
        parts = {}

        try:
            for matrix_part in ('data', 'indices', 'indptr', 'shape'):
                parts[matrix_part] = getattr(f.root.matrix, matrix_part).read()

        except Exception as e:
            return None  # Return None instead of an incomplete matrix

        matrix = csr_matrix((parts['data'], parts['indices'], parts['indptr']), shape=parts['shape'])

        f.close()  # Manually close file

    return matrix

import glob
import os

# Find all .hic files in the directory
hic_files = glob.glob(os.path.join(data_dir, "*chr2_10kb.h5"))

for file in hic_files:
    print(file)
    matrix = read_h5(file)
    # processing steps..

解决方案与分析 #

1. 核心问题定位 #

Colab的共享运行时环境确实可能存在HDF5缓存残留问题，尤其是tables库处理大文件时，即便关闭文件句柄，底层HDF5缓存可能未被及时释放。另外，你的代码中f.close()属于冗余操作——with上下文管理器会自动关闭文件，手动调用反而可能引发异常，建议删除该行。

2. 强制内存释放优化 #

在每次迭代后，除了删除变量，还需清空tables的全局文件缓存，并强制触发垃圾回收：

import gc
import tables
from psutil import Process

for file in hic_files:
    print(file)
    matrix = read_h5(file)
    # 执行你的处理步骤
    # --- 内存清理步骤 ---
    del matrix
    # 清空tables所有未关闭的文件句柄
    tables.file._open_files.close_all()
    # 强制垃圾回收
    gc.collect()
    # 触发内存统计，帮助Colab释放闲置后台缓存
    Process().memory_info()

3. 分块读取避免内存峰值 #

如果文件过大，一次性读取整个CSR矩阵会瞬间占用大量内存，建议分块读取data/indices，降低单次内存占用：

import numpy as np
from scipy.sparse import csr_matrix

def read_h5_chunked(file_name, chunk_size=10**6):
    with tables.open_file(file_name, 'r') as f:
        shape = getattr(f.root.matrix, 'shape').read()
        indptr = getattr(f.root.matrix, 'indptr').read()
        
        data_list = []
        indices_list = []
        # 按行分段读取数据
        for i in range(len(indptr)-1):
            start, end = indptr[i], indptr[i+1]
            if end - start > chunk_size:
                # 大段数据再分块读取
                for chunk_start in range(start, end, chunk_size):
                    chunk_end = min(chunk_start + chunk_size, end)
                    data_list.append(f.root.matrix.data[chunk_start:chunk_end].read())
                    indices_list.append(f.root.matrix.indices[chunk_start:chunk_end].read())
            else:
                data_list.append(f.root.matrix.data[start:end].read())
                indices_list.append(f.root.matrix.indices[start:end].read())
        
        data = np.concatenate(data_list)
        indices = np.concatenate(indices_list)
        matrix = csr_matrix((data, indices, indptr), shape=shape)
    return matrix

4. Colab运行时特殊处理 #

• 避免在循环内重复导入库（将import glob/import os等放在循环外）
• 若问题持续，重启Colab运行时，清除之前的缓存残留
• 切换到高内存运行时（Colab设置→更改运行时类型→内存规格选择高内存）

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