相关疑难解决方法(0)

我有两个函数的实现,它计算frobenius范数并减去跟踪.此函数应用于4D张量x的维3中的所有向量.然后总结所有结果.我用它作为一个回忆录的一部分.TensorFlow的版本是0.9.

我的第一个实现使用了tf.batch_*函数.

def test1(x):
    """x: [batch, height, width, channels]"""
    s = x.get_shape().as_list()
    a = tf.reshape(x, [-1, s[3], 1])
    c = tf.batch_matmul(a, a, adj_y=True)
    c2 = tf.square(c)
    diag = tf.batch_matrix_diag_part(c2)
    return tf.reduce_sum(c2) - tf.reduce_sum(diag)

这是有效的,但中间张量c是大于张量x的通道时间,这限制了我的批量大小.因此,我尝试了一种基于map_fn的方法:

def fn(x):
    x1 = tf.reshape(x, [-1, 1])
    c1 = tf.matmul(x1, x1, transpose_b=True)
    c2 = tf.square(c1)
    t1 = tf.trace(c2)
    return tf.reduce_sum(c2)- t1)

def test2(x):
    """x: [batch, height, width, channels]"""
    s = x.get_shape().as_list()
    a = tf.reshape(x, [-1, s[3]])
    return tf.reduce_sum(tf.map_fn(fn, a))

当我运行第二个函数时,我收到许多(50+)个消息,如:

I tensorflow/core/common_runtime/gpu/pool_allocator.cc:244] 
PoolAllocator: After 16084 get …