Оптимизировать использование памяти функцией Python ⇐ Python

[Anonymous]

Обновлено
У меня есть следующая функция, которая используется в пространственной статистической модели. Я могу запустить модель на подмножестве моих данных (1 штат США), используя несколько ГБ ОЗУ, но когда я расширяю ее на всю территорию США, ей не хватает памяти.
Я внес несколько изменений в исходную функцию. Насколько это возможно, я использую разреженные представления и методы, разработанные для разреженных матриц. Входная матрица A имеет размер 207972 x 207972 и содержит 1314488 ненулевых элементов, или около 0,003%. При тестировании полного набора данных в кластере и использовании 500 ГБ ОЗУ (пробовано до 750 ГБ) код выполняется до Sigma = spsolve(Q_perturbed, b).astype(np.float32), прежде чем закончится память. Я могу сэкономить память, используя int16 и т. д., потому что большинство элементов являются одними, но я понимаю, что это приведет к увеличению накладных расходов, поскольку Python преобразуется в float32, чтобы гарантировать, что все имеет одинаковый тип данных при выполнении некоторых операций.
Мне нужно запустить функцию только один раз, чтобы получить число, а затем я могу использовать это число в качестве фиксированных входных данных для моей модели. Он застревает в spsolve, несмотря на то, что функция использует разреженные матрицы. Насколько я понимаю, эта функция может генерировать плотные промежуточные матрицы. Я попробовал представления как csc, так и csr — csr рекомендовался как более эффективный. Я также сейчас запускаю процесс для каждого столбца в цикле, используя 60 ГБ ОЗУ, чтобы избежать проблем с ОЗУ. Однако он очень медленный - т. е. не завершает 10 000 столбцов (около 5%) примерно за 20 минут. Я также попробовал использовать гораздо больший объем оперативной памяти (1 ТБ) без циклического метода, и мне не хватило памяти. Кластер может обрабатывать до 2 ТБ оперативной памяти, но я думаю, что для запуска потребуется неделя. Я рассмотрел несколько других решений, которые заменяют решатель spsolve().
from scipy.sparse import diags, identity, csc_matrix
from scipy.linalg import solve
from scipy.sparse.linalg import spsolve
def blockwise_solve(Q_perturbed, b):
"""
Solve the system Q_perturbed * x = b using block-wise solving,
where each column of b represents a separate right-hand side.
"""
n = Q_perturbed.shape[0]
solutions = []

# Iterate over columns of b (each column is a separate right-hand side)
for i in range(b.shape[1]):
if i%10**4==0:
logging.info(i)
print(i)
rhs = b[:, i].toarray().flatten() # Convert column to 1D array
solution = spsolve(Q_perturbed, rhs, use_umfpack=True) # Solve for this right-hand side
solutions.append(solution)

# Combine solutions into a matrix
return np.vstack(solutions).T # Return the solutions as a matrix, one solution per column

def scaling_factor_sp(A):
"""Compute the scaling factor from an adjacency matrix.
This function uses sparse matrix computations and is most
efficient on sparse adjacency matrices. Used in the BYM2 model.
The scaling factor is a measure of the variance in the number of
edges across nodes in a connected graph.
Only works for fully connected graphs. The argument for scaling
factors is developed by Andrea Riebler, Sigrunn H. Sørbye,
Daniel Simpson, Havard Rue in "An intuitive Bayesian spatial
model for disease mapping that accounts for scaling"
https://arxiv.org/abs/1601.01180"""

# Computes the precision matrix in sparse format.
num_neighbors = A.sum(axis=1).A.ravel().astype(np.float32)
D = diags(num_neighbors, format="csc", dtype=np.float32) # Degree matrix
del num_neighbors
Q = D - A # Precision matrix
del D
del A

# add a small jitter along the diagonal
jitter = max(Q.diagonal()) * np.sqrt(np.finfo(np.float32).eps)
Q_perturbed = Q + diags(np.ones(Q.shape[0]) * jitter, dtype=np.float32, format="csc")
del jitter
del Q
# Compute a version of the pseudo-inverse

n = Q_perturbed.shape[0]
# b = identity(n, dtype=np.int8, format="csc")
b = identity(n, format="csc").astype(np.float32)
gc.collect()
Q_perturbed = Q_perturbed.tocsr()
# Solve the system using block-wise solving
Sigma = blockwise_solve(Q_perturbed, b)

del Q_perturbed
del b
gc.collect()

W = Sigma.dot(np.ones(Sigma.shape[0], dtype=np.float32))
Q_inv = Sigma - (W[:, np.newaxis] * W[np.newaxis, :]) / W.sum()
del W
del Sigma
gc.collect()

# Compute the geometric mean of the diagonal on a
# precision matrix.

return np.exp(np.sum(np.log(Q_inv.diagonal())) / Q_inv.shape[0])

scaling_factor = scaling_factor_sp(adj_matrix)

Вывод с использованием полного набора данных.
2024-11-08 10:47:23,524 - INFO - After running Q Memory Usage: 171.34 MB
2024-11-08 10:47:23,552 - INFO - After running Q_perturbed Memory Usage: 179.16 MB
2024-11-08 10:47:23,553 - INFO - After running b Memory Usage: 179.16 MB
2024-11-08 10:47:23,553 - INFO - Sparsity of Q_perturbed: 0.0035%


Подробнее здесь: https://stackoverflow.com/questions/791 ... n-function