Keras 3 层性能基准测试全指南：用 benchmarks/layer_benchmark 对比 keras.layers 与 tf.keras.layers 的前向与训练性能-平芜编程栈

人工智能
深度学习
机器学习
预训练

【免费下载链接】keras

Deep Learning for humans

项目地址：https://gitcode.com/GitHub_Trending/ke/keras

本指南围绕 Keras 仓库中 benchmarks/layer_benchmark 目录展开，系统讲解如何通过命令行基准脚本对比同一层（Layer）在 Keras 3 与 TensorFlow 内置 tf.keras 两套 API 下的执行性能。读完本文，你将掌握基准脚本的全部命令行参数、前向传播与训练步两种测速模式的工作原理、各基准文件内置的层与配置清单，以及如何自定义基准层来测量自己的模型组件。

基准套件定位：同一层，两套 API，双向对比

benchmarks/layer_benchmark目录下的脚本用于比较keras.layers.XXX与tf.keras.layers.XXX两个同名层的性能差异。它不只测前向传播（forward pass），还会测训练步（train step，即前向 + 反向传播）的整体吞吐量。这两套 API 虽然同名，但分属 Keras 3（独立的多后端实现）与 TensorFlow 内置的旧版 Keras，执行路径与算子调度存在差异，因此基准脚本用完全相同的层配置、输入形状与数据生成逻辑，在同一个进程中先后运行两套模型，保证对比条件一致。

每个基准模块都是一个独立可执行的 Python 模块（基于absl的app/flags框架），通过python3 -m直接运行。完整的模块清单见 benchmarks/layer_benchmark 下的各文件。

快速上手：运行第一个基准

以卷积层为例，在仓库根目录执行：

python3 -m benchmarks.layer_benchmark.conv_benchmark \ --benchmark_name=benchmark_conv2D \ --num_samples=2048 \ --batch_size=256 \ --jit_compile=True

运行后，脚本会依次打印两个阶段的结果：

前向传播（forward pass）吞吐量：Keras 3 throughput of forward pass of Conv2D: xxx samples/sec.与TF Keras throughput of forward pass of Conv2D: xxx samples/sec.
训练步（forward & backward pass）吞吐量：同样以 samples/sec 为单位分别打印 Keras 3 与 TF Keras 的结果。

需要注意，benchmark_name必须与目标模块BENCHMARK_NAMES字典中注册的函数名严格一致，否则会抛出ValueError（如 conv_benchmark.py 所示，错误信息会列出所有合法的名字）。如果省略--benchmark_name，则该模块注册的所有基准会按顺序全部执行一遍（对应 conv_benchmark.py 的main分支）。

命令行参数详解

所有基准模块共享同一套参数定义，位于 base_benchmark.py，通过absl.flags注册：

参数	类型	默认值	含义
`--benchmark_name`	string	`None`	要运行的基准名称；为`None`时运行该模块注册的全部基准
`--num_samples`	integer	`1000`	输入数据样本数
`--batch_size`	integer	`20`	数据批大小
`--jit_compile`	bool	`True`	为`True`时以 XLA 编译方式运行

三个数值参数在 conv_benchmark.py 中被读取并原样传给对应的基准函数：

benchmark_name = FLAGS.benchmark_name num_samples = FLAGS.num_samples batch_size = FLAGS.batch_size jit_compile = FLAGS.jit_compile

关于jit_compile需要说明的是：它被同时传递给 Keras 3 模型与 tf.keras 模型（见 base_benchmark.py），两者都会以loss="mse"、optimizer="sgd"编译，仅在jit_compile开关上保持一致，从而在 XLA 开/关两种模式下分别对比两套实现。

样本数与批大小的搭配建议：从源码看，测速迭代次数由num_iterations = num_samples // batch_size - 1决定（base_benchmark.py），即去掉首个 batch（首批常包含 XLA 编译等预热开销）后的完整批次数。因此建议让num_samples至少为batch_size的 2 倍以上，否则测速区间可能为空，无法得到有效吞吐量。

两种测速模式：前向传播与训练步

每个基准函数都会先后调用benchmark_predict与benchmark_train（例如 conv_benchmark.py），分别覆盖：

前向传播（predict 模式）：benchmark_predict(num_samples, batch_size, data=None)（base_benchmark.py）。在未显式传入data时，会按输入形状自动生成高斯随机数据np.random.normal(size=[num_samples] + list(input_shape))；对于多输入层则逐输入生成。随后以相同batch_size先后调用keras_model.predict(...)与tf_keras_model.predict(...)，各自挂载测速回调，输出 samples/sec。

训练步（fit 模式）：benchmark_train(num_samples, batch_size, data=None, label=None)（base_benchmark.py）。标签默认由“层对随机数据的前向输出乘以 1.001”生成，注释明确说明这是为了“避免零梯度”（Scale by a small factor to avoid zero gradients.），从而保证 SGD 训练过程有效。随后通过model.fit(data, label, batch_size=..., callbacks=[...])完成一个 epoch 的训练，测得的吞吐量涵盖前向与反向传播全流程。

这两条路径共用同一套吞吐量统计机制（见下文），因此两个阶段的结果可直接横向比较 Keras 3 与 TF Keras 的差距。

基准框架源码剖析：LayerBenchmark 与测速回调

所有基准模块都构建在 base_benchmark.py 之上，核心是LayerBenchmark类（base_benchmark.py）：

class LayerBenchmark: def __init__( self, layer_name, init_args, input_shape, flat_call_inputs=True, jit_compile=True, keras_layer=None, tf_keras_layer=None, ):

其关键设计如下：

1. 同名层的双实现解析：layer_name只是一个字符串，构造时通过getattr(keras.layers, layer_name)与getattr(tf.keras.layers, layer_name)动态获取两套层的类（base_benchmark.py），再用同一份init_args实例化。如果keras_layer/tf_keras_layer参数非空，则直接使用调用方传入的预构造层——这是为Bidirectional这类需要分别传入keras.layers.Layer与tf.keras.layers.Layer的复合层预留的扩展点（base_benchmark.py 中的注释明确说明了这一用途）。

2. 函数式模型的统一搭建：_build_keras_model与_build_tf_keras_model（base_benchmark.py）会先把input_shape归一化为列表（非列表/元组时包成单元素列表），逐形状创建keras.Input/tf.keras.Input，然后根据flat_call_inputs决定调用方式：

flat_call_inputs=True：outputs = self._keras_layer(*inputs)，即把多个输入按位置参数展开传给层（如MultiHeadAttention、Dot的部分用法）；
flat_call_inputs=False：outputs = self._keras_layer(inputs)，即把输入列表作为一个整体参数传入（如Attention、AdditiveAttention、Concatenate等多输入层）。

最终封装为keras.Model/tf.keras.Model，并统一以mse损失与sgd优化器编译。

3. 吞吐量统计的回调机制：测速不依赖模型自带的时间统计，而是通过一个自定义回调完成。BenchmarkMetricsCallback 记录从start_batch（默认第 1 个 batch）到stop_batch之间的耗时，吞吐量公式为：

throughput = (stop_batch - start_batch + 1) / (end_time - begin_time)

即“测速区间内的批次数 ÷ 耗时”，单位是 batch/s；再乘上batch_size换算为 samples/sec（base_benchmark.py）。KerasCoreBenchmarkMetricsCallback与TFKerasBenchmarkMetricsCallback（base_benchmark.py）分别把它包装为keras.callbacks.Callback与tf.keras.callbacks.Callback，从而可以挂载到两套模型的predict/fit调用上。由于从第 1 个 batch 才开始计时，首批的编译/预热开销被排除在统计之外，结果更能反映稳定运行时的吞吐。

内置基准矩阵：各模块覆盖的层与默认配置

目录下每个基准模块覆盖一类层，以下配置均取自各模块源码中的init_args与input_shape，可直接对照复现：

卷积层（conv_benchmark.py）：Conv1D（filters=64, kernel_size=2, 输入[1024, 256]）、Conv2D（filters=16, kernel_size=2, 输入[128, 128, 4]）、Conv3D（filters=16, kernel_size=2, 输入[32, 32, 32, 4]）、DepthwiseConv1D/2D（kernel_size=16, depth_multiplier=2）、SeparableConv1D/2D（kernel_size=16, depth_multiplier=2, filters=3）、Conv1DTranspose（filters=32, kernel_size=4）、Conv2DTranspose 与 Conv3DTranspose（filters=16, kernel_size=2）。共 10 个基准。

激活层（activation_benchmark.py）：ELU、ReLU、LeakyReLU、PReLU、Softmax，均使用默认参数，输入[256, 256]。

注意力层（attention_benchmark.py）：Attention（双输入[[256, 64], [256, 64]]，flat_call_inputs=False）、MultiHeadAttention（num_heads=4, key_dim=16，三输入[[256, 64], [256, 64], [256, 64]]）、AdditiveAttention（三输入，flat_call_inputs=False）。

核心层（core_benchmark.py）：Dense（units=256，输入[256, 256]）、EinsumDense（equation="abc,cd->abd"，output_shape=(None, 256)）、Embedding（input_dim=128, output_dim=256，输入[256]）。其中 Embedding 的输入数据为整数索引，因此该基准显式传入data = [np.random.randint(30, size=(num_samples, 256))]（core_benchmark.py），这也演示了“层需要特殊数据分布时如何覆盖默认随机数据”的用法。

合并层（merge_benchmark.py）：Add、Average、Concatenate、Maximum、Minimum、Subtract（双输入[[256, 256], [256, 256]]）、Dot（axes=[2, 1]，输入[[256, 32], [32, 64]]）、Multiply（输入[[256, 64], [256, 64]]），全部使用flat_call_inputs=False。

归一化层（normalization_benchmark.py）：BatchNormalization（输入[256, 256, 4]）、GroupNormalization（groups=2，输入[256, 256, 4]）、LayerNormalization（输入[256, 128, 4]）、UnitNormalization（输入[256, 128, 4]）。

池化层（pooling_benchmark.py）：AveragePooling1D/2D/3D 与 MaxPooling1D/2D/3D（pool_size=2），以及 GlobalAveragePooling1D/2D/3D 与 GlobalMaxPooling1D/2D/3D，共 12 个基准，输入维度分别对应 1D[1024, 256]、2D[256, 256, 3]、3D[64, 64, 32, 3]。

正则化层（regularization_benchmark.py）：Dropout（rate=0.5，输入[256, 256, 4]）、GaussianDropout（rate=0.5）、GaussianNoise（stddev=0.5）、SpatialDropout1D/2D/3D（rate=0.5，输入分别[256, 3]、[256, 256, 3]、[32, 32, 32, 3]）。

重塑层（reshaping_benchmark.py）：Cropping1D/2D/3D、Flatten、Permute、UpSampling1D/2D/3D、ZeroPadding1D/2D/3D，共 11 个基准。

循环层（rnn_benchmark.py）：ConvLSTM1D/2D/3D、GRU、LSTM、SimpleRNN（输入[256, 256]）、Bidirectional（输入[256, 256]）、TimeDistributed（输入[10, 32, 32, 3]）。

图像增强层（random_rotation_benchmark.py）：RandomRotation（factor=0.1，输入[224, 224, 3]），模拟典型图像输入尺寸。

每个模块的BENCHMARK_NAMES字典将函数名映射到基准函数，--benchmark_name即由此解析（如 pooling_benchmark.py 注册了全部 12 个池化基准）。

进阶用法：基准自己的层

LayerBenchmark的构造参数已经覆盖了大多数自定义需求，无需改动框架代码即可扩展新的基准：

多输入层：将input_shape传入嵌套列表，并用flat_call_inputs控制调用约定。例如注意力基准中Attention用flat_call_inputs=False（整体传列表），而MultiHeadAttention用flat_call_inputs=True（展开为位置参数），原因在于两层的call签名不同——前者接收一个inputs列表，后者接收query/value/key多个参数。

需要特殊初始化或复合结构的层：通过keras_layer/tf_keras_layer直接传入预构造实例。以Bidirectional为例，它需要包装一个具体的循环层，且该循环层在 Keras 3 与 TF Keras 下分别实例化；框架的注释明确指出这是该参数存在的意义（base_benchmark.py）。

非高斯数据分布：benchmark_predict与benchmark_train都接受data/label参数。Embedding 基准就是典型例子——它需要整数索引输入，因此显式传入np.random.randint生成的离散数据，避免默认的高斯连续数据破坏层的语义。

运行环境与依赖

基准脚本的对比对象是tf.keras.layers，因此运行环境必须安装 TensorFlow（base_benchmark.py 顶部直接import tensorflow as tf）。仓库根目录提供了requirements.txt、requirements-tensorflow-cuda.txt、requirements-tensorflow-tpu.txt等依赖清单，可按硬件环境选用。此外脚本基于absl（app与flags），同样需要预先安装。执行时务必以仓库根目录作为工作目录、以-m模块方式运行，才能保证benchmarks.layer_benchmark.*包导入路径可用。

结果解读与使用建议

单位为 samples/sec：两个阶段（predict / fit）的输出都是吞吐量而非耗时，数值越大越快；由于两套实现使用完全相同的输入数据、批大小与优化器，结果差异可直接归因于层实现本身。
区分前向与训练步：前向吞吐反映单次推理路径；训练步吞吐叠加了反向传播与优化器开销，更接近真实训练负载。若迁移或优化层实现，两个指标应分别关注。
预热已剔除：回调从第 1 个 batch 才开始计时，首批的 XLA 编译等一次性开销不计入结果，因此多次运行的结果稳定性较好。
控制变量：修改--num_samples/--batch_size会影响内存占用与吞吐绝对值，跨配置对比时保持二者一致，--jit_compile开关也应固定。