Блог

Я определил свою функцию потерь и хочу выполнить итерацию каждого элемента пакета, чтобы вычислить функцию потерь. Я использовал tf.map_fn , однако, я обнаружил, что это очень медленно. Есть ли какие-нибудь предложения?

 def loss(phi, mu, sigma, t_phi, t_mu, t_sigma):
    _loss = 0.0
    for i in range(phi.shape[0]):
        for j in range(phi.shape[0]):
            _loss  = phi[i] * phi[j] * pdf(mu[i], mu[j], tf.sqrt(sigma[i]**2   sigma[j]**2))
            _loss  = t_phi[i] * t_phi[j] * pdf(t_mu[i], t_mu[j], tf.sqrt(t_sigma[i]**2   t_sigma[j]**2))
            _loss  = -2 * phi[i] * t_phi[j] * pdf(mu[i], t_mu[j], tf.sqrt(sigma[i]**2   t_sigma[j]**2))
    return tf.sqrt(_loss)

def reduce_loss(phi, mu, sigma, t_phi, t_mu, t_sigma):
    with tf.variable_scope('loss') as loss:
        stacked = tf.stack([phi, mu, sigma, t_phi, t_mu, t_sigma], 1)
        return tf.map_fn(lambda x: loss(x[0], x[1], x[2], x[3], x[4], x[5]), stacked,
                         parallel_iterations=4)

def pdf(x, mu, sigma):
    return tf.exp(-0.5*(x-mu)**2/sigma**2) / ((2*np.pi*sigma**2)**0.5)
 

Размер пакета равен 1024.

Вы можете устранить циклы в своей loss функции. Это делается путем векторизации всего. Например, вы перебираете i и j вычисляете phi[i]*phi[j] , но это ij-й элемент tf.matmul(phi[:, None], phi[None, :]) . Выполнение этого должно быть быстрее, чем реализация с циклами.

Кроме того, поскольку tensorflow строит график статически, вашей функции может потребоваться много времени, чтобы даже построить график. По этой причине вам обычно следует избегать больших вложенных циклов for в tensorflow.

Я привел пример с векторизованной частью вашей функции потерь, остальные части должно быть легко выполнить.

 import tensorflow as tf
from numpy import pi as PI
from time import time


# some random vectors
size = 10
phi = tf.random.uniform([size])
mu = tf.random.uniform([size])
sigma = tf.random.uniform([size])


####################################
# Your original loss
####################################

def pdf(x, m, s):
    return tf.exp(-0.5*(x-m)**2/s**2) / ((2*PI*s**2)**0.5)


def loss():
    _loss = 0.0
    for i in range(phi.shape[0]):
        for j in range(phi.shape[0]):
            _loss  = phi[i] * phi[j] * pdf(mu[i], mu[j], tf.sqrt(sigma[i]**2   sigma[j]**2))
    return tf.sqrt(_loss)


####################################
# vectorised loss
####################################

def vector_pdf(x, s):
    return tf.exp(-0.5*x**2/s**2) / ((2*PI*s**2)**0.5)


def vectorised_loss():
    phi_ij = tf.matmul(phi[:, None], phi[None, :])
    difference = mu[:, None] - mu[None, :]
    sigma_squared = sigma**2
    sigma_sum = tf.sqrt(sigma_squared[:, None]   sigma_squared[None, :])

    loss_array = phi_ij*vector_pdf(difference, sigma_sum)
    return tf.sqrt(tf.reduce_sum(loss_array))


#######################################
# Time the functions and show they are the same
#######################################

with tf.Session() as sess:
    loop_loss = loss()
    vector_loss = vectorised_loss()
    # init = tf.global_variables_initializer()
    # sess.run(init)

    t = 0.
    for _ in range(100):
        st = time()
        loop_loss_val = sess.run(loop_loss)
        t  = time() - st
    print('loop took {}'.format(t/100))

    t = 0.
    for _ in range(100):
        st = time()
        vector_val = sess.run(vector_loss)
        t  = time() - st
    print('vector took {}'.format(t / 100))

    l_val, v_val = sess.run([loop_loss, vector_loss])
    print(l_val, v_val)
 

Это печатает

 loop took 0.01740453243255615
vector took 0.004280190467834472
4.6466274 4.6466274
 

Векторизируя функцию потерь, ваша функция уменьшения также должна быть простой для векторизации. Теперь вам нужно будет выполнить пакетную обработку matmul и немного изменить индексы вычитаний. Например:

 mu[:, None] - mu[None, :]
# becomes
mu[: ,:, None] - mu[:, None, :]