CODEWITHSUNDEEP
pythontensorflow2.0
Igor Rivin
Igor Rivin

Reputation: 4864

tensorflow performance confusion

I am using tensorflow for a non-neural-net project (interacting particles), and have two version of the potential function. Here is version 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

I thought that breaking out individual x, y, z was kind of ugly, and probably slow, so here is a slicker version 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

The good news is that the two versions return exactly the same results. The bad news is that the slick version is much, much slower (almost a factor of two). Does anyone understand why this would be, or what the "canonical" way to write such code is?

Upvotes: 0

Views: 148

Answers (1)

guest
guest

Reputation: 11

This isn't really good enough for an answer, but I also don't have enough cred to comment...so I'd suggest maybe not using tf.norm, and instead compute the norm by hand like you did in the first Coulomb potential function.

(I know this isn't about numpy, but you can also compare np.linalg.norm against math.sqrt(np.sum(x*x))...and on my systems I see the former is significantly slower. Maybe both np and tf compute the norm using singular values?)

Upvotes: 1

Related Questions