如何在TensorFlow中获取Transformer计算图并计算其FLOPs？

你可以根据你使用的TensorFlow版本，选择对应方式获取Transformer的计算图：

TensorFlow 1.x 静态图模式 #

直接在定义Transformer前向逻辑的上下文管理器中获取当前静态图即可，示例代码如下：

import tensorflow as tf
# 导入你自定义的Transformer实现，也可以用官方内置实现
from your_model_file import Transformer

graph = tf.Graph()
with graph.as_default():
    # 定义输入占位符，形状根据你的实际场景调整
    input_seq = tf.placeholder(tf.int32, shape=[batch_size, seq_length])
    target_seq = tf.placeholder(tf.int32, shape=[batch_size, seq_length])
    # 初始化Transformer并执行前向传播
    transformer = Transformer(vocab_size=32000, d_model=512, num_layers=6)
    output = transformer(input_seq, target_seq, training=True)
    
    # 直接传入当前graph统计FLOPs
    flops = tf.profiler.profile(
        graph, 
        options=tf.profiler.ProfileOptionBuilder.float_operation()
    )
    print(f"Transformer总FLOPs：{flops.total_float_ops}")

TensorFlow 2.x 动态图模式 #

2.x默认开启Eager执行，没有默认静态图，需要用tf.function装饰Transformer的前向调用过程，触发静态图追踪后再获取计算图，示例代码如下：

import tensorflow as tf
from your_model_file import Transformer

# 初始化Transformer
transformer = Transformer(vocab_size=32000, d_model=512, num_layers=6)
# 构造示例输入，形状根据你的实际场景调整
batch_size = 1
seq_length = 512
input_seq = tf.random.uniform((batch_size, seq_length), maxval=32000, dtype=tf.int32)
target_seq = tf.random.uniform((batch_size, seq_length), maxval=32000, dtype=tf.int32)

# 用tf.function装饰前向逻辑，触发静态图追踪
@tf.function
def forward_step(input_seq, target_seq):
    return transformer(input_seq, target_seq, training=False)

# 调用一次完成静态图构建
_ = forward_step(input_seq, target_seq)
# 获取追踪完成的静态图
graph = forward_step.get_concrete_function(input_seq, target_seq).graph

# 统计FLOPs，高版本TF如果找不到tf.profiler，可以替换为tf.compat.v1.profiler.profile
flops = tf.profiler.profile(
    graph, 
    options=tf.profiler.ProfileOptionBuilder.float_operation()
)
print(f"Transformer总FLOPs：{flops.total_float_ops}")

补充说明 #

如果你使用的是Keras官方内置的tf.keras.layers.Transformer层，上述逻辑完全适用，只需要替换掉自定义Transformer的初始化逻辑即可。

内容的提问来源于stack exchange，提问作者Maxwell Albert