本文翻译自： 《Prototyping kernels and advanced visualization with Python ops》， 如有侵权请联系删除， 仅限于学术交流， 请勿商用。 如有谬误， 请联系指出。
TensorFlow 中的内核操作完全用 C ++编写， 以提高效率。 但是用 C++编写 TensorFlow 内核的话可能会非常痛苦。 因此， 在花费数小时实现属于自己的内核之前， 你也许需要先实现一个操作的原型， 尽管这样的效率会很低。 通过 tf.py_func() 你可以将任何一个 python 源代码转换为 TensorFlow 的操作。
举个例子而言， 这里有一个用 python 自己实现的 ReLU 非线性激活函数， 通过 tf.py_func() 转换为 TensorFlow 操作的例子：
import numpy as np
import tensorflow as tf
import uuid
def relu(inputs):
# Define the op in python
def _relu(x):
return np.maximum(x, 0.)
# Define the op's gradient in python
def _relu_grad(x):
return np.float32(x > 0)
# An adapter that defines a gradient op compatible with TensorFlow
def _relu_grad_op(op, grad):
x = op.inputs[0]
x_grad = grad * tf.py_func(_relu_grad, [x], tf.float32)
return x_grad
# Register the gradient with a unique id
grad_name = "MyReluGrad_" + str(uuid.uuid4())
tf.RegisterGradient(grad_name)(_relu_grad_op)
# Override the gradient of the custom op
g = tf.get_default_graph()
with g.gradient_override_map({"PyFunc": grad_name}):
output = tf.py_func(_relu, [inputs], tf.float32)
return output
通过 TensorFlow 的gradient checker， 你可以确认这些梯度是否计算正确：
x = tf.random_normal([10])
y = relu(x * x)
with tf.Session():
diff = tf.test.compute_gradient_error(x, [10], y, [10])
print(diff)
compute_gradient_error() 数值化地计算梯度， 返回与理论上的梯度的差别， 我们所期望的是一个非常小的差别。
注意到我们的这种实现是非常低效率的， 这仅仅在实现模型原型的时候起作用， 因为 python 代码并不能并行化而且不能在 GPU 上运算（导致速度很慢）。 一旦你确定了你的 idea， 你就需要用 C++重写其内核。
在实践中， 我们一般在 Tensorboard 中用 python 操作进行可视化。 如果你是在构建一个图片分类模型， 而且想要在训练过程中可视化你的模型预测， 那么 TF 允许你通过 tf.summary.image() 函数进行图片的可视化。
image = tf.placeholder(tf.float32)
tf.summary.image("image", image)
但是这仅仅是可视化了输入的图片， 为了可视化其预测结果， 你还必须找一个方法在图片上添加预测标识， 当然这在现有的 tensorflow 操作中是不存在的。 一个更简单的方法就是通过 python 将预测标志绘制到图片上， 然后再封装它。
import io
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
def visualize_labeled_images(images, labels, max_outputs=3, name="image"):
def _visualize_image(image, label):
# Do the actual drawing in python
fig = plt.figure(figsize=(3, 3), dpi=80)
ax = fig.add_subplot(111)
ax.imshow(image[::-1,...])
ax.text(0, 0, str(label),
horizontalalignment="left",
verticalalignment="top")
fig.canvas.draw()
# Write the plot as a memory file.
buf = io.BytesIO()
data = fig.savefig(buf, format="png")
buf.seek(0)
# Read the image and convert to numpy array
img = PIL.Image.open(buf)
return np.array(img.getdata()).reshape(img.size[0], img.size[1], -1)
def _visualize_images(images, labels):
# Only display the given number of examples in the batch
outputs = []
for i in range(max_outputs):
output = _visualize_image(images[i], labels[i])
outputs.append(output)
return np.array(outputs, dtype=np.uint8)
# Run the python op.
figs = tf.py_func(_visualize_images, [images, labels], tf.uint8)
return tf.summary.image(name, figs)
请注意， 因为 summary 通常只评估一次（并不是每步都执行）， 因此可以在实践中可以使用而不必担心效率。