python - Tensorflow 张量中的自定义函数

Question

假设我想通过以下方式在afunction上应用：tensorflow tensor

|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 speedup。

标签： pythontensorflow

Answer 1

我想我可以自己回答这个问题，我的函数如下所示：

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)