Блог

Я ищу наилучший способ вычислить сферические функции Бесселя с возрастающим порядком вдоль одной оси тензора. Мой тензор имеет форму (n, x, y) , где n обозначает порядок функции Бесселя. Для этого кода я упростил их, j_n(x) = x**n но для реальных функций сложность возрастает очень быстро для больших n .

Мне нужно только вычислить верхний треугольник моего тензора, поэтому количество точек уменьшается с увеличением значений n. Это может выглядеть так:

 import tensorflow as tf
import numpy as np

test_data = tf.random.normal((40, 40, 100)) 
x = tf.where(np.triu(np.ones((40, 40)))[..., None] , test_data, tf.zeros_like(test_data))
x[:, :, 0]

<tf.Tensor: shape=(40, 40), dtype=float32, numpy=
array([[-1.157, -1.376, -0.427, ...,  0.031,  0.846,  0.813],
       [ 0.   , -2.165,  0.873, ..., -1.044,  0.999,  0.778],
       [ 0.   ,  0.   ,  2.528, ...,  0.409,  0.041, -1.473],
       ...,
       [ 0.   ,  0.   ,  0.   , ...,  1.587,  1.987, -0.626],
       [ 0.   ,  0.   ,  0.   , ...,  0.   , -1.137, -0.859],
       [ 0.   ,  0.   ,  0.   , ...,  0.   ,  0.   , -0.901]],
      dtype=float32)>

      x_sparse = tf.sparse.from_dense(x)
 

На данный момент у меня есть два подхода к решению этой проблемы: оба используют tf.map_fn либо полный тензор, который требует много ненужных вычислений, либо разреженный тензор, который даже для больших n работает медленнее.

 def func(inp):
    x, a = inp
    return x**a

def func_sparse(inp):
    x, a = inp
    return tf.sparse.SparseTensor(x.indices, x.values**a, x.dense_shape)

%timeit tf.map_fn(func, (x, tf.range(40, dtype=tf.float32)), fn_output_signature=tf.float32)
# 18.8 ms ± 89.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit tf.map_fn(func_sparse, (x_sparse, tf.range(40, dtype=tf.float32)), fn_output_signature=tf.SparseTensorSpec([None, None], dtype=tf.float32))
# 30.1 ms ± 166 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
 

Есть ли лучший способ применить функцию с возрастающим порядком вдоль оси тензора? %timeit Вышесказанное не учитывает генерацию разреженного тензора, что делает его еще медленнее. Есть ли способ избежать ненужных вычислений и при этом быть быстрее и эффективнее с точки зрения памяти?