tensorflow - 在 tf 2.x 中以图形模式运行 TensorFlow op
问题描述
我想对一些 TensorFlow 操作进行基准测试(例如在它们之间或针对 PyTorch)。但是大多数时候我会写一些类似的东西:
import numpy as np
import tensorflow as tf
tf_device = '/GPU:0'
a = np.random.normal(scale=100, size=shape).astype(np.int64)
b = np.array(7).astype(np.int64)
with tf.device(tf_device):
a_tf = tf.constant(a)
b_tf = tf.constant(b)
%timeit tf.math.floormod(a_tf, b_tf)
这种方法的问题在于它在急切模式下进行计算(我认为特别是它必须执行 GPU 到 CPU 的放置)。最终,我想在tf.keras
模型中使用这些操作,因此想评估它们在图形模式下的性能。
首选的方法是什么?
我的谷歌搜索没有任何结果,我不知道如何使用 tf 1.x 中的会话。
解决方案
您正在寻找的是tf.function
. 检查本教程和本文档。
正如教程所说,在 TensorFlow 2 中,Eager Execution 默认开启。用户界面直观且灵活(运行一次性操作更容易、更快捷),但这可能会以牺牲性能和可部署性为代价。要获得高性能和可移植的模型,请使用 tf.function 从您的程序中制作图表。
检查此代码:
import numpy as np
import tensorflow as tf
import timeit
tf_device = '/GPU:0'
shape = [100000]
a = np.random.normal(scale=100, size=shape).astype(np.int64)
b = np.array(7).astype(np.int64)
@tf.function
def experiment(a_tf, b_tf):
tf.math.floormod(a_tf, b_tf)
with tf.device(tf_device):
a_tf = tf.constant(a)
b_tf = tf.constant(b)
# warm up
experiment(a_tf, b_tf)
print("In graph mode:", timeit.timeit(lambda: experiment(a_tf, b_tf), number=10))
print("In eager mode:", timeit.timeit(lambda: tf.math.floormod(a_tf, b_tf), number=10))
推荐阅读
- typescript - Typescript:使用联合类型和三元运算符进行类型推断
- kubernetes - Kubernetes DNS 使用 systemd resolvd 的理想设置
- html - 如何在视频背景上覆盖导航栏
- charts - chartjs默认背景颜色到列
- jquery - 使用 jQuery CSS 在 Botframework 中更改自适应卡片
- java - 重构 for 循环和错误(类型不匹配:无法从元素类型 String 转换为 Integer)
- .htaccess - 从类别内的 url 中删除/签名
- c# - 识别字符串内容,并验证
- amazon-s3 - 在单键和多键分区雅典娜中哪个性能更高?
- python - 我应该在运行新程序时编辑 VS Code 的环境吗