Python Tensorflow张量中的自定义函数

假设我想通过以下方式对

张量流张量

应用

函数

：

|2 7 4|      |f(2) f(7) f(4)|
|6 7 2| ---> |f(6) f(7) f(2)| ---> with f(x)=2+e^x
|6 8 8|      |f(6) f(8) f(8)|

---

在tensorflow层中，在GPU支持下，如何实现并行化以获得所需的速度。使用

tf.py_func

的方法很简单，使用numpy，但速度很慢。我更喜欢一个完全基于tensorflow的解决方案-->

tf GPU加速

我想我自己可以回答这个问题，我的功能如下：

Function of interest: f(a1,a2,a3,x)=a1+a2*e^(x/a3)

相应的tensorflow代码：

import tensorflow as tf

# Input data
img = tf.reshape(tf.range(5*5*1, dtype=tf.float32), [1, 5, 5, 1])

a1=2.5
a2=10
a3=75

img=tf.math.divide(img,a3)
img=tf.math.exp(img)
img=tf.math.multiply(img,a2)
img=tf.math.add(img,a1)

检查结果：

print(img)