Keras 3 架构解析（一）：基础计算层——从计算后端到张量运算

引言

深度学习框架的核心任务，是将数学模型转换为能够在计算设备上执行的程序。一个完整的深度学习系统不仅需要描述神经网络的结构，还必须解决张量运算、自动微分、计算图编译、内存管理以及硬件执行等问题。

在传统的深度学习框架中，这些能力通常被整合在同一套技术体系内。例如，PyTorch 具有自己的张量系统、自动微分引擎与计算编译机制；TensorFlow 则提供张量运算、计算图构建及相应的运行时系统。

Keras 3 采用了不同的架构设计。它不再将模型开发绑定到某一个计算框架，而是通过统一的 API，将上层模型定义与底层计算实现分离，使同一套模型代码能够运行在 JAX、TensorFlow 和 PyTorch 等不同后端之上。

这一设计使 Keras 形成了两个相互关联的层次：

• 框架抽象层：负责提供统一的张量运算接口、模型组件以及训练 API。
• 计算执行层：由选定的后端负责自动微分、程序编译、设备执行等底层功能。

理解这种分层关系，是理解 Keras 3 整体架构的基础。

本文围绕 Keras 的基础计算层展开，重点讨论计算后端、Keras Ops、自动微分、JIT 编译、随机数管理以及数值精度与设备管理等核心机制。

一、计算后端：Keras 的执行基础

1.1 什么是计算后端

计算后端（Backend）是负责执行实际数值计算的软件系统。

在深度学习任务中，矩阵乘法、卷积、归一化和非线性激活等操作最终都必须转化为计算设备可以执行的指令。计算后端负责将这些高层数学操作映射到对应的计算实现，并管理执行过程所需要的张量和计算资源。

Keras 3 本身并不重新实现一套完整的底层计算引擎，而是将这些工作委托给已经成熟的计算框架。

其主要计算后端包括：

三者都可以完成神经网络训练，但其内部运行机制存在显著差异。

Keras 的作用，是在这些差异之上建立相对统一的开发接口。

需要特别指出的是，多后端兼容不意味着不同后端具有相同的执行性能，也不保证完全一致的数值结果。

不同框架采用的算子实现、编译策略和浮点计算顺序可能存在差异。因此，即使运行完全相同的 Keras 模型，也可能出现训练速度、设备利用率或最终数值结果的差别。

1.2 后端的选择与配置

Keras 3 可以通过环境变量 KERAS_BACKEND 选择计算后端。

例如，使用 JAX：

import os

os.environ["KERAS_BACKEND"] = "jax"

import keras

print(keras.config.backend())

将 "jax" 替换为 "tensorflow" 或 "torch"，即可选择对应的计算后端。

需要注意，Keras 和对应的后端框架必须预先安装。设置环境变量并不负责安装所需依赖。

此外，后端应当在导入 Keras 之前完成配置。虽然 Keras 也提供运行时切换后端的相关接口，但这可能导致已有对象失效，因此在实际开发中，通常更适合通过进程启动配置确定后端，并在切换时重新启动 Python 进程。

这种设计使模型代码与计算环境之间形成了相对清晰的边界。

例如，同一个项目可以根据硬件和实际需求，分别在使用 JAX 的 TPU 环境和使用 PyTorch 的 GPU 环境中运行，而不必分别维护两套完全独立的模型定义。

但这一能力存在前提：模型必须使用与后端无关的 Keras API，不能在核心计算逻辑中依赖某个后端独有的操作。

二、Keras Ops：统一的张量运算接口

2.1 为什么需要统一算子接口

不同深度学习框架虽然支持类似的数学运算，但它们的 API 并不完全相同。

例如，矩阵乘法可以分别写成：

# TensorFlow
tf.matmul(a, b)

# PyTorch
torch.matmul(a, b)

# JAX
jax.numpy.matmul(a, b)

这些代码表达了相同的数学意图，但依赖不同框架的张量类型及运算实现。

如果直接使用某个框架的原生 API 编写自定义神经网络层，那么该层通常也会受到该框架的限制。

为解决这一问题，Keras 3 提供了统一的张量运算接口：keras.ops。

上述矩阵乘法可以统一写成：

from keras import ops

result = ops.matmul(a, b)

Keras 将根据当前选择的后端，调用相应的底层计算实现。

开发者不需要为了不同后端而重复编写同一段数学运算逻辑。

2.2 Keras Ops 的主要组成

keras.ops 不仅覆盖基本矩阵运算，还提供了范围较广的计算功能。

第一类：基础数学与张量操作

包括：

• ops.add()：加法
• ops.multiply()：乘法
• ops.matmul()：矩阵乘法
• ops.mean()：求均值
• ops.sum()：求和
• ops.reshape()：张量形状变换
• ops.transpose()：张量维度转置
• ops.concatenate()：张量拼接

这些接口构成神经网络数值计算的基础，并在设计上广泛借鉴 NumPy 的 API 风格。

第二类：神经网络运算

包括：

• ops.relu()：ReLU 激活函数
• ops.softmax()：Softmax 函数
• ops.sigmoid()：Sigmoid 函数
• ops.conv()：卷积运算
• ops.binary_crossentropy()：二元交叉熵

这一类算子直接服务于深度学习模型构建。

第三类：线性代数及其他计算

包括矩阵分解、范数计算、傅里叶变换、图像处理以及部分控制流操作等。

对于 Transformer、CNN 和其他复杂神经网络而言，这些功能可以在不直接调用后端原生 API 的情况下完成大量计算逻辑。

2.3 张量与符号表示

Keras 的张量操作需要区分两种不同的使用情境。

第一种是实际数值计算。此时参与运算的是当前后端能够处理的真实张量。

例如：

from keras import ops

a = ops.ones((2, 3))
b = ops.ones((3, 4))

c = ops.matmul(a, b)

print(ops.shape(c))
# (2, 4)

第二种是符号式模型构建。

Keras 提供 KerasTensor，用于描述模型构建阶段的张量形状、数据类型及符号计算关系。

例如：

import keras
from keras import ops

inputs = keras.Input(shape=(8,))

outputs = ops.relu(inputs)

model = keras.Model(inputs, outputs)

这里的 inputs 并不是一批已经载入的实际训练数据，而是一个用于表示模型输入的符号对象。

Keras 可以据此建立模型的计算关系，后续再通过实际数据执行计算。

因此，KerasTensor 与 JAX、TensorFlow 或 PyTorch 的原生张量并不是完全相同的概念。前者主要承担 Keras 符号建模中的结构描述职责，后者承担具体的数值计算职责。

在 Functional API 中，应使用支持 KerasTensor 的 Keras 操作构建计算关系，而不是任意将其传入后端原生函数。

2.4 跨后端自定义层

统一算子接口的重要价值，体现在自定义神经网络层的开发中。

下面实现一个简单的线性变换层：

import keras
from keras import ops

class Linear(keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.kernel = self.add_weight(
            name="kernel",
            shape=(input_shape[-1], self.units),
            initializer="glorot_uniform",
        )

        self.bias = self.add_weight(
            name="bias",
            shape=(self.units,),
            initializer="zeros",
        )

    def call(self, inputs):
        x = ops.matmul(inputs, self.kernel)
        return ops.relu(x + self.bias)

该层的前向计算对应：

[
Y=\operatorname{ReLU}(XW+b)
]

其中，(X) 是输入矩阵，(W) 是权重矩阵，(b) 是偏置向量。

由于计算过程使用了 keras.ops，而参数通过 Keras 提供的权重管理接口创建，因此该层可以在兼容的 Keras 后端中运行。

这也是 Keras 3 最重要的设计原则之一：模型组件定义应尽可能独立于具体计算后端。

不过，这种抽象并不能消除所有底层差异。当模型需要使用特定 GPU 内核、自定义 TPU 算子或后端独有功能时，开发者仍可能需要直接访问后端原生接口。

三、自动微分：神经网络训练的数学基础

3.1 自动微分的基本原理

神经网络训练的基本目标，是通过调整模型参数，使损失函数逐渐减小。

假设模型参数为 (\theta)，损失函数为：

[
L=f(\theta)
]

采用梯度下降方法时，参数更新规则为：

[
\theta_{t+1}

\theta_t-\eta\nabla_{\theta}L
]

其中，(\eta) 为学习率，(\nabla_{\theta}L) 为损失函数相对于模型参数的梯度。

自动微分（Automatic Differentiation）用于高效计算这些梯度。

它不同于数值微分，也不同于直接进行复杂表达式的符号求导。自动微分通过记录或变换程序中的基本运算，并利用链式法则计算导数。

深度学习中的反向传播，本质上可以视为反向模式自动微分在神经网络中的典型应用。

3.2 Keras 如何实现自动微分

严格而言，Keras 并不为所有后端重新实现一套完全独立的自动微分引擎。

Keras 负责管理模型变量、损失函数、优化器及训练过程，并调用当前后端的自动微分能力。

例如：

JAX

JAX 通过函数变换实现自动微分，典型接口为：

jax.grad()
jax.value_and_grad()

TensorFlow

TensorFlow 通常通过 tf.GradientTape 记录运算并计算梯度。

PyTorch

PyTorch 使用 Autograd 系统，通过计算图追踪张量运算，并利用反向传播计算梯度。

这些机制的实现方式有所不同，但都可以为 Keras 模型训练提供必要的梯度信息。

在使用 Keras 内置训练接口时，开发者通常不需要直接编写后端特定的梯度计算过程。

例如：

model.compile(
    optimizer="adam",
    loss="mse"
)

model.fit(
    x_train,
    y_train,
    epochs=10
)

Keras 会在训练过程中协调前向计算、损失计算、梯度计算和参数更新。

不过，如果需要实现完全自定义的底层训练循环，仍可能需要调用 JAX、TensorFlow 或 PyTorch 各自的自动微分 API。

这体现了 Keras 的另一项重要架构特征：统一的高层训练接口并不意味着底层梯度计算过程完全相同。

四、计算图与 JIT 编译

4.1 即时执行与编译执行

深度学习程序通常涉及大量重复的数学运算。如果每次执行都需要由 Python 逐项调度算子，就可能产生额外开销。

为了提高效率，现代深度学习框架普遍提供了计算图及编译优化能力。

常见的执行方式可以概括为两类。

即时执行（Eager Execution）

程序按操作顺序执行，适合交互式开发、调试和快速验证。

编译执行（Compiled Execution）

系统首先分析程序中的计算关系，再通过编译器生成经过优化的执行程序。

编译过程可能包括：

• 消除冗余计算
• 融合多个计算算子
• 优化内存访问
• 调整运算顺序
• 生成适合目标硬件的计算代码

这类优化对于大规模神经网络训练尤其重要。

4.2 不同后端的编译机制

Keras 的多后端设计，并没有统一底层编译器。

在 JAX 后端中，常见方式是使用 jax.jit 结合 XLA 编译机制，对符合要求的函数进行追踪与编译。

TensorFlow 可以通过 tf.function 等机制构建计算图，并在适用情况下使用 XLA 优化。

PyTorch 则提供 torch.compile，其实现涉及图捕获、图变换及后端代码生成等机制。

这些编译系统并不等价。它们支持的算子范围、图捕获行为和动态形状处理能力可能存在差异。

4.3 Keras 对 JIT 的支持

Keras 在模型编译接口中提供 jit_compile 参数：

model.compile(
    optimizer="adam",
    loss="mse",
    jit_compile="auto"
)

该参数控制是否尝试对训练过程启用相应的编译优化。

根据 Keras 的 API 文档，在 JAX 和 TensorFlow 后端中，"auto" 会在模型支持的情况下启用 XLA 编译；在 PyTorch 后端中，"auto" 默认采用即时执行，而显式设置 True 可以使用基于 torch.compile 的编译路径。

但应当注意，JIT 编译并不必然提高程序性能。

首先，编译过程本身需要消耗时间。对于执行次数较少的任务，编译开销可能超过运行时收益。

其次，输入形状变化、复杂控制流或不受支持的算子，可能导致重复编译、编译失败或性能下降。

因此，在性能优化时，应同时关注首次编译耗时、稳定运行吞吐量以及设备利用率，而不能仅以是否启用了 JIT 作为评价标准。

五、随机数生成与可复现性

随机性是深度学习训练的重要组成部分。

神经网络的权重初始化、Dropout、数据增强及随机采样等操作，都需要使用随机数。

但在 JAX、TensorFlow 和 PyTorch 中，随机数管理存在不同的机制。

为了提供相对统一的行为，Keras 3 提供了 keras.random 接口。

5.1 无状态随机数接口

Keras 的随机数操作采用显式种子控制。

例如：

import keras

a = keras.random.normal(
    shape=(2, 3),
    seed=42
)

b = keras.random.normal(
    shape=(2, 3),
    seed=42
)

对于相同的后端和计算环境，使用相同整数种子重复调用，通常会得到相同的随机结果。

但是，如果希望在多次调用中生成不同的随机数，就不能简单地反复传入同一个整数种子。

5.2 SeedGenerator

Keras 提供 SeedGenerator 管理连续随机数生成所需要的种子状态。

import keras

seed_gen = keras.random.SeedGenerator(42)

a = keras.random.normal(
    shape=(2, 3),
    seed=seed_gen
)

b = keras.random.normal(
    shape=(2, 3),
    seed=seed_gen
)

上述两次调用会使用不同的随机状态。

这对于需要在训练过程中不断生成随机掩码的 Dropout 等操作尤为重要。

特别是在 JAX 后端中，JIT 追踪对隐式全局随机状态存在限制，因此编写需要兼容 JAX 编译的自定义层时，应正确使用局部 SeedGenerator。

需要区分的是，随机数可复现并不等于整个训练过程完全可复现。

实际训练还可能受到并行计算、非确定性 GPU 内核、浮点运算顺序及分布式执行等因素的影响。

因此，设置随机种子只是实现实验可复现性的必要措施之一。

六、数据类型与计算设备管理

6.1 数值精度

深度学习模型通常采用浮点数表示权重、激活值及梯度。

常见数据类型包括：

降低计算精度可以减少部分数据的内存占用，并利用现代加速器的低精度计算能力提高吞吐量。

但低精度同样可能引入数值稳定性问题。

例如，FP16 的指数范围相对较窄，训练过程中可能出现数值上溢或下溢。BF16 则保留了与 FP32 相近的指数范围，但有效尾数精度较低。

两者适用于不同的硬件和计算场景。

6.2 混合精度策略

Keras 通过 DTypePolicy 将计算精度与变量存储精度分开管理。

例如：

import keras

keras.config.set_dtype_policy(
    "mixed_bfloat16"
)

layer = keras.layers.Dense(128)

print(layer.compute_dtype)
print(layer.variable_dtype)

对于该混合精度策略，通常可以得到：

bfloat16
float32

这意味着模型层可以使用 BF16 执行主要浮点计算，同时将变量保存在 FP32 精度下。

对于 GPU，也可以根据硬件支持情况选择：

keras.config.set_dtype_policy(
    "mixed_float16"
)

混合精度并不意味着所有算子都会无条件采用低精度。实际数值计算还受到算子实现、后端行为以及数值稳定性要求的影响。

特别需要注意，设置全局浮点类型与设置混合精度策略并不相同。直接把全部计算都设置为 FP16，可能使训练更容易出现数值问题，因此一般应优先使用明确的混合精度策略。

6.3 计算设备管理

神经网络最终需要运行在 CPU、GPU 或 TPU 等硬件上。

Keras 可以通过选定的后端使用相应的计算设备，但设备发现、底层内存管理、编译器支持和内核执行主要由后端及其运行时负责。

例如，在 JAX 后端中，可以通过：

import jax

print(jax.devices())

查看 JAX 识别的计算设备。

对于多设备训练，Keras 还提供了分布式计算接口，例如设备网格与张量布局抽象。不过，统一分布式 API 的后端支持范围与各框架自身的分布式能力并不完全相同，需要根据具体版本和执行环境进行确认。

需要强调的是，Keras 的精度策略与计算设备选择属于不同层次的概念。

前者主要决定张量采用何种数值表示及计算精度；后者决定计算如何映射到物理硬件。

两者共同影响训练性能，但不能相互替代。

七、基础计算层的架构意义

通过上述分析，可以将 Keras 3 的基础计算架构概括为三个层次。

第一层：Keras 统一接口

包括 keras.ops、keras.random、数值类型策略以及统一的模型与训练抽象。

这一层负责向开发者提供相对稳定的编程接口。

第二层：后端计算框架

包括 JAX、TensorFlow 和 PyTorch。

它们负责提供实际张量表示、自动微分、编译与运行时能力。

第三层：硬件执行系统

包括 CPU、GPU、TPU，以及相应的底层计算库、编译器和设备运行环境。

模型的实际性能最终取决于整个计算链路的协同工作，而不仅仅取决于上层接口。

从软件工程的角度看，Keras 3 的设计体现了接口与实现分离的原则。

模型定义被放置在相对独立的抽象层中，计算后端则成为可替换的执行实现。

这种结构的主要价值不在于消除不同计算框架之间的差异，而在于将差异限制在明确的架构边界内，使模型开发者能够在较少修改业务代码的情况下选择不同的计算系统。

当然，抽象也存在成本。统一接口无法覆盖所有后端特性，也不能保证任何硬件上的最优性能。在对算子融合、内存布局、编译图结构或设备利用率有严格要求的场景中，仍可能需要深入后端实现，甚至编写硬件专用算子。

因此，Keras 更适合被理解为建立在成熟计算框架之上的统一深度学习开发体系，而不是这些框架的完全替代品。

结语

Keras 3 基础计算层的核心，是通过统一的计算接口，将模型的数学定义与实际计算实现解耦。

这一层主要解决三个问题：

1. 如何让相同的模型运算在不同计算后端上执行。
2. 如何统一管理张量运算、随机数、数据精度等基础计算行为。
3. 如何在保持上层接口稳定的同时，利用各后端提供的自动微分、JIT 编译及硬件计算能力。

掌握这些机制后，便可以更准确地理解 Keras 中 Layer、Model 和训练系统的职责边界。

其中，基础计算层负责定义和执行基本数学操作，神经网络层进一步组织这些操作并管理可训练参数，模型层则将多个网络层组合成完整的计算结构。

在下一部分中，我们将进一步讨论 Keras 的神经网络层系统，分析 Layer 基类、权重与状态管理、build() 和 call() 的生命周期，以及神经网络组件如何建立在基础计算层之上。