基于大量大型字典构建大型稀疏矩阵的优化方案咨询

问题背景 #

我有一系列格式为dict_k:{(i,j):d_ij}的字典，其中i、j是待构建稀疏矩阵的整数索引，d_ij为浮点数，每个字典最多包含100万键值对。共有约160个这类字典，每个大小约16MB，字典间可能存在键值重复（如不同字典中都有(0,0):1.25）。需要构建包含所有键值对的稀疏矩阵，最终用于执行矩阵向量乘法，矩阵规模约为500万×500万。

初始方案是构建巨型字典再转COO矩阵，但过程导致电脑卡顿：

bigDict={}
for i in range(160):
    with open(dictionary_path, "rb") as fp:   
        bigDict.update(pickle.load(fp))

优化方案 #

1. 跳过巨型字典，直接收集数据构建矩阵 #

构建巨型字典会产生极高的内存峰值（仅字典结构本身就会占用额外内存），直接拆分每个字典的键值到三个列表（行索引、列索引、值）是更高效的方式，同时可按需处理重复键：

场景A：重复键的值需要累加

这是矩阵构建的常规需求，直接收集所有条目后转COO再转CSR（CSR是矩阵向量乘法效率最高的格式），SciPy会自动合并相同(i,j)位置的值：

import numpy as np
from scipy.sparse import coo_matrix, csr_matrix

rows = []
cols = []
vals = []

# 遍历所有字典文件
for idx in range(160):
    dict_path = f"your_dict_{idx}.pkl"
    with open(dict_path, "rb") as fp:
        current_dict = pickle.load(fp)
        # 拆分键值对到列表
        for (r, c), v in current_dict.items():
            rows.append(r)
            cols.append(c)
            vals.append(v)

# 构建COO矩阵并转换为CSR
coo = coo_matrix((vals, (rows, cols)), shape=(5_000_000, 5_000_000))
csr_matrix = coo.tocsr()

# 执行矩阵向量乘法
result = csr_matrix.dot(your_vector)

场景B：重复键保留最后一次出现的值

如果需要用后读取字典的值覆盖之前的重复键，需先去重再构建矩阵：

import numpy as np
from scipy.sparse import csr_matrix

rows = []
cols = []
vals = []

for idx in range(160):
    dict_path = f"your_dict_{idx}.pkl"
    with open(dict_path, "rb") as fp:
        current_dict = pickle.load(fp)
        rows.extend(r for r, c in current_dict.keys())
        cols.extend(c for r, c in current_dict.keys())
        vals.extend(current_dict.values())

# 转换为numpy数组便于处理
rows_np = np.array(rows, dtype=np.int64)
cols_np = np.array(cols, dtype=np.int64)
vals_np = np.array(vals, dtype=np.float64)

# 按行、列排序，保留最后一个重复项
sorted_indices = np.lexsort((cols_np, rows_np))
sorted_rows = rows_np[sorted_indices]
sorted_cols = cols_np[sorted_indices]
sorted_vals = vals_np[sorted_indices]

# 生成唯一掩码，筛选最后出现的条目
unique_mask = np.concatenate([[True], (sorted_rows[1:] != sorted_rows[:-1]) | (sorted_cols[1:] != sorted_cols[:-1])])
final_rows = sorted_rows[unique_mask]
final_cols = sorted_cols[unique_mask]
final_vals = sorted_vals[unique_mask]

# 构建CSR矩阵
csr_matrix = csr_matrix((final_vals, (final_rows, final_cols)), shape=(5_000_000, 5_000_000))

2. 分块处理降低内存峰值 #

如果内存仍紧张，可分批次处理字典，每批构建一个小CSR矩阵后合并：

from scipy.sparse import csr_matrix
import pickle

total_csr = None
block_size = 10  # 每批处理10个字典

for idx in range(0, 160, block_size):
    batch_rows = []
    batch_cols = []
    batch_vals = []
    # 处理当前批次的字典
    for batch_idx in range(idx, min(idx + block_size, 160)):
        dict_path = f"your_dict_{batch_idx}.pkl"
        with open(dict_path, "rb") as fp:
            current_dict = pickle.load(fp)
            for (r, c), v in current_dict.items():
                batch_rows.append(r)
                batch_cols.append(c)
                batch_vals.append(v)
    # 构建批次CSR矩阵
    batch_csr = csr_matrix((batch_vals, (batch_rows, batch_cols)), shape=(5_000_000, 5_000_000))
    # 合并到总矩阵（累加场景）
    if total_csr is None:
        total_csr = batch_csr
    else:
        total_csr += batch_csr

3. 替换pickle为numpy二进制存储（长期优化） #

如果可以重新生成字典文件，建议直接存储为numpy数组（.npy），避免pickle的序列化开销，读取速度更快：

生成文件时：

import numpy as np

# 假设current_dict是当前要存储的字典
rows = np.array([r for r, c in current_dict.keys()], dtype=np.int64)
cols = np.array([c for r, c in current_dict.keys()], dtype=np.int64)
vals = np.array(list(current_dict.values()), dtype=np.float64)

np.save(f"rows_{idx}.npy", rows)
np.save(f"cols_{idx}.npy", cols)
np.save(f"vals_{idx}.npy", vals)

读取构建矩阵时：

import numpy as np
from scipy.sparse import csr_matrix

rows_list = []
cols_list = []
vals_list = []

for idx in range(160):
    rows_list.append(np.load(f"rows_{idx}.npy"))
    cols_list.append(np.load(f"cols_{idx}.npy"))
    vals_list.append(np.load(f"vals_{idx}.npy"))

# 合并所有数组
rows_np = np.concatenate(rows_list)
cols_np = np.concatenate(cols_list)
vals_np = np.concatenate(vals_list)

csr_matrix = csr_matrix((vals_np, (rows_np, cols_np)), shape=(5_000_000, 5_000_000))

关键注意事项 #

• 矩阵向量乘法优先使用CSR格式，其在稀疏矩阵-稠密向量乘法上的效率远高于COO、CSC等格式。
• 500万×500万的矩阵需确保索引值不超出范围，避免构建矩阵时出错。

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