python之Tensorflow，迭代张量

我定义了我的损失函数，我想迭代批处理中的每个项目来计算损失函数。我使用了 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]但这是 tf.matmul(phi[:, None], phi[None, :]) 的第 ij 个元素.这样做应该比使用循环实现更快。

此外，由于 tensorflow 静态构建图形，您的函数甚至可能需要很长时间来构建图形。出于这个原因，您通常应该避免在 tensorflow 中使用大的嵌套 for 循环。

我已经将你的损失函数的一部分向量化了，其他部分应该很容易做。

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 

通过矢量化损失函数，您的 reduce 函数也应该可以直接矢量化。现在您需要批量处理 matmul，并稍微更改减法的索引。例如:

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