Повышение скорости производительности при пакетном вычислении расстояния Махаланобиса. ⇐ Python

[Anonymous]

У меня есть следующий фрагмент кода, который вычисляет расстояние Махаланобиса по набору пакетных функций. На моем устройстве это занимает около 100 мс, большая часть которого связана с матричным умножением между delta и inv_covariance
дельта — это матрица размером 874x32x100, inv_covariance — размерностью 874x100x100

Код: Выделить всё

def compute_distance(embedding: np.ndarray, mean: np.ndarray, inv_covariance: np.ndarray) -> np.ndarray:
batch, channel, height, width = embedding.shape
embedding = embedding.reshape(batch, channel, height * width)

# calculate mahalanobis distances
delta = np.ascontiguousarray((embedding - mean).transpose(2, 0, 1))

distances = ((delta @ inv_covariance) * delta).sum(2).transpose(1, 0)
distances = distances.reshape(batch, 1, height, width)
distances = np.sqrt(distances.clip(0))

return distances

Я пытался преобразовать код для использования numba и @njit, я заранее выделил промежуточную матрицу и пытаюсь выполнить умножение меньших матриц с помощью цикла for, поскольку matmul не поддерживается для трехмерных матриц.

Код: Выделить всё

def compute_distance(embedding: np.ndarray, mean: np.ndarray, inv_covariance: np.ndarray) -> np.ndarray:
batch, channel, height, width = embedding.shape
embedding = embedding.reshape(batch, channel, height * width)

# calculate mahalanobis distances
delta = np.ascontiguousarray((embedding - mean).transpose(2, 0, 1))
inv_covariance = np.ascontiguousarray(inv_covariance)

intermediate_matrix = np.zeros_like(delta)
for i in range(intermediate_matrix.shape[0]):
intermediate_matrix[i] = delta[i] @ inv_covariance[i]

distances = (intermediate_matrix * delta).sum(2).transpose(1, 0)
distances = np.ascontiguousarray(distances)
distances = distances.reshape(batch, 1, height, width)
distances = np.sqrt(distances.clip(0))

return distances

Я добавил несколько последовательных массивов, последний важен, иначе код не работает, остальные добавлены для подавления предупреждения о том, что @ будет работать быстрее (похоже, это не так). слишком много).
Есть ли способ ускорить код, улучшив его или переосмыслив его другим математическим способом?
Править - Окончательная реализация
На основе ответа Жерома Ришара у меня получился этот код

Код: Выделить всё

@nb.njit()
def matmul(delta: np.ndarray, inv_covariance: np.ndarray):
"""Computes distances = ((delta[i] @ inv_covariance[i]) * delta[i]).sum(2) using numba.

Args:
delta: Matrix of dimension BxD
inv_covariance: Matrix of dimension DxD

Returns:
Matrix of dimension BxD
"""
si, sj, sk = delta.shape[0], inv_covariance.shape[1], delta.shape[1]
assert sk == inv_covariance.shape[0]
line = np.zeros(sj, dtype=delta.dtype)
res = np.zeros(si, dtype=delta.dtype)
for i in range(si):
line.fill(0.0)
for k in range(sk):
factor = delta[i, k]
for j in range(sj):
line[j] += factor * inv_covariance[k, j]
for j in range(sj):
res[i] += line[j] * delta[i, j]
return res

@nb.njit
def mean_subtraction(embeddings: np.ndarray, mean: np.ndarray):
"""Computes embeddings - mean using numba, this is required as I have errors with the default numpy
implementation.

Args:
embeddings: Embedding matrix of dimension FxBxD
mean: Mean matrix of dimension BxD

Returns:
Delta matrix of dimension FxBxD
"""
output_matrix = np.zeros_like(embeddings)
for i in range(embeddings.shape[0]):
output_matrix[i] = embeddings[i] - mean

return output_matrix

@nb.njit(parallel=True)
def compute_distance_numba(embedding: np.ndarray, mean: np.ndarray, inv_covariance: np.ndarray) -> np.ndarray:
"""Compute distance score using numba.

Args:
embedding: Embedding Vector
mean: Mean of the multivariate Gaussian distribution
inv_covariance: Inverse Covariance matrix of the multivariate Gaussian distribution.
"""
batch, channel, height, width = embedding.shape
embedding = embedding.reshape(batch, channel, height * width)

delta = np.ascontiguousarray(mean_subtraction(embedding, mean).transpose(2, 0, 1))
inv_covariance = np.ascontiguousarray(inv_covariance)

intermediate_matrix = np.zeros((delta.shape[0], delta.shape[1]), dtype=delta.dtype)
for i in nb.prange(intermediate_matrix.shape[0]):
intermediate_matrix[i] = matmul(delta[i], inv_covariance[i])

distances = intermediate_matrix.transpose(1, 0)
distances = np.ascontiguousarray(distances)
distances = distances.reshape(batch, 1, height, width)
distances = np.sqrt(distances.clip(0))

return distances

Изменения по сравнению с принятым ответом — это пользовательская функция вычитания и добавление dtype для промежуточной матрицы, чтобы избежать использования np.float64 по умолчанию.

Подробнее здесь: https://stackoverflow.com/questions/774 ... omputation