Блог

Я использую 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
 

Хорошей новостью является то, что две версии дают совершенно одинаковые результаты. Плохая новость заключается в том, что скользкая версия намного, намного медленнее (почти в два раза). Кто-нибудь понимает, почему это было бы, или каков «канонический» способ написания такого кода?

Это на самом деле недостаточно хорошо для ответа, но у меня также недостаточно репутации, чтобы comment…so Я бы предложил, возможно, не использовать tf.norm, а вместо этого вычислить норму вручную, как вы сделали в первой функции кулоновского потенциала.

(Я знаю, что речь идет не о numpy, но вы также можете сравнить np.linalg.norm с math.sqrt(np.sum(x*x)) …и в моих системах я вижу, что первый значительно медленнее. Может быть, и np, и tf вычисляют норму, используя сингулярные значения?)