如何在张量流中对批处理维度上的函数进行成对评估？

【问题标题】：How to evaluate a function pair-wise over batch dimension in tensorflow?如何在张量流中对批处理维度上的函数进行成对评估？
【发布时间】：2019-12-12 15:31:25
【问题描述】：

给定形状为 [a, n] 的张量 x 和形状为 [b,n] 的 y 以及函数 f(p,q) 对形状为 [1,n]（或 [n]）的张量进行运算并返回一个标量值，我怎样才能在 x 和 y 的批量维度上成对计算 f 以便我得到的张量是 [a,b,1]（或 [a,b]）？

我知道这适用于乘法和加法等运算，如下所述：Evaluate all pair combinations of rows of two tensors in tensorflow 通过隐式广播。

如何将其扩展到任意函数？

应用程序是我想计算两个张量的成对 KL 散度以匹配它们，所以基本上是蛮力 NN 计算。

【问题讨论】：

标签： python tensorflow nearest-neighbor

【解决方案1】：

一种解决方案（我可以使用 TF 最好的方法）是使用两个 map_fn，如下所示。这会很慢，因为map_fn 没有内在的并行化。但是你应该能够通过修改parallel_iteration 参数来获得一些加速。

import tensorflow as tf

a = tf.random_normal(shape=[10,5])
b = tf.random_normal(shape=[8,5])

def f(x,y):
  return x+y

res = tf.map_fn(lambda x: tf.map_fn(lambda y: f(x,y),b), a)

with tf.Session() as sess:
  print(sess.run(res))

【讨论】：

使用我要评估的函数，由于某种原因，这会立即使 CUDA 求解器崩溃，但这是一个幼稚的解决方案。