Question

我的代码使用了以下两个功能，因此我希望尽可能提高效率。

def batch_inner_broadcast(density, core, opt_einsum=False):
    '''compute M_k=A_k^T*L*A_k'''
    if opt_einsum:
        M  = tf.einsum('bsu,ksr->kurb', density, core)
        return tf.einsum('kurb,kut->bkrt', M, core)
    else:
        return tf.einsum('krs,bsu,kut->bkrt', tf.transpose(core, [0,2,1]), density, core)

def batch_inner_contraction(density, core, weights = None, opt_einsum=False):
    '''compute Sum_k w_k*A_k^T*L*A_k'''
    if opt_einsum:
        if weights is not None:
            M = tf.einsum('bsu,ksr->kurb', density, core)
            M = tf.einsum('kurb,kut->krtb', M, core)
            return tf.einsum('krtb,bk->brt', M, weights)
        else:
            M = tf.einsum('kut,ksr->surt', core, core)
            return tf.einsum('surt,bsu->brt', M, density)
    else:
        if weights is not None:
            return tf.einsum('bkij,bk->bij', batch_inner_broadcast(density, core), weights)
        else:
            return tf.einsum('krs,bsu,kut', tf.transpose(core, [0,2,1]), density, core)

我认为Tensorflow的einsum可能不是最佳的，所以我尝试从opt_einsum包中提取最佳减少路径，但它们似乎没什么帮助。似乎Numpy优化可能无法很好地转移到Tensorflow。这些函数在opt_einsum=True时使用这些假定的更好的例程。

i,j,s,r,u,t指数的大小可能不同。目前它们很小，但如果它很好的话会很棒。 b是批量维度，在100-1000范围内。 k通常很小。

编辑：

的小时间测试

core = tf.random_normal((5,10,3))
density = tf.random_normal((5000, 10, 10))
weight = tf.random_normal((5000,5))

bib_no = batch_inner_broadcast(density, core, opt_einsum=False)
bic_no = batch_inner_contraction(density, core, weight, opt_einsum=False)
bib_yes = batch_inner_broadcast(density, core, opt_einsum=True)
bic_yes = batch_inner_contraction(density, core, weight, opt_einsum=True)

%timeit sess.run(bib_no)
%timeit sess.run(bic_no)
%timeit sess.run(bib_yes)
%timeit sess.run(bic_yes)

产量

25.6 ms ± 769 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
39.8 ms ± 717 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
36.2 ms ± 300 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
40.3 ms ± 738 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

编辑：更高级的问题涉及嵌套的batch_inner_contraction，其中结果被输入到下一个的密度字段中。

Tensorflow：加速多维线性代数

0 个答案: