张量流性能混乱

Question

我将tensorflow用于一个非神经网络项目（相互作用粒子），并且有两个版本的势函数。以下是第1版：

def coulomb(t, fudge = tf.constant(1000000.0)):
  howmany = t.shape[0]
  x = t[:, 0:1]
  y = t[:, 1:2]
  z = t[:, 2:3]

  dx = tf.transpose(x) - x
  dy = tf.transpose(y) - y
  dz = tf.transpose(z) - z
  rmat = tf.sqrt(dx **2  + dy **2 + dz **2)
  rmat = tf.linalg.set_diag(tf.cast(rmat, tf.float32), fudge * tf.ones(howmany))

  inv = 1.0/rmat

  PE = tf.reduce_sum(inv)/2

  return PE

我认为爆发个人x, y, z有点难看，而且可能很慢，所以这里有一个更圆滑的版本2：

def coulomb2(t, fudge = tf.constant(1000000.0)):
  howmany = t.shape[0]

  tt = tf.expand_dims(t, 1)
  difft = tf.transpose(tt, perm=(1, 0, 2)) - tt
  rmat = tf.norm(difft, axis=2)
  rmat = tf.linalg.set_diag(tf.cast(rmat, tf.float32), fudge * tf.ones(howmany))

  inv = 1.0/rmat

  PE = tf.reduce_sum(inv)/2

  return PE

好消息是这两个版本返回的结果完全相同。坏消息是，光滑的版本要慢得多（几乎是两倍）。有人知道为什么会这样，或者写这样代码的“规范”方式是什么吗

Answer 1

这对答案来说还不够好，但我也没有足够的可信度来评论…所以我建议不要使用tf.norm，而是像第一个库仑势函数那样手工计算范数

（我知道这与numpy无关，但您也可以将np.linalg.norm与math.sqrt(np.sum(x*x))进行比较……在我的系统上，我发现前者的速度要慢得多。可能np和tf都使用奇异值计算范数？）