本文深入解析 TensorFlow 中

tf.Module

tf.Module

tf.keras.layers.Layer

tf.keras.Model

接下来将按照以下结构逐步剖析：核心定位 → 基础代码逐行解读 → 关键特性详解 → 复杂模型实战示例，涵盖类结构设计、关键 API 使用、变量收集机制等细节内容。

一、

tf.Module

的核心作用与定位

tf.Module

tf.Variable

• 自动变量收集：所有赋值给实例属性的

  tf.Variable

  会被自动归类为可训练或全部变量，无需手动维护列表；
• 支持嵌套模块组合：一个

  tf.Module

  实例可以包含多个其他

  tf.Module

  子模块（如模型中包含多个层），并递归地汇总所有层级的变量；
• 适配模型生命周期操作：为后续的训练流程（梯度更新）、模型保存与加载（仅需存储变量数值）提供基础支撑。

形象比喻来说，

tf.Module

tf.Module

二、第一个示例解析：简易模块 SimpleModule

该示例展示了

tf.Module

y = a*x + b

1. 类定义与初始化（继承 tf.Module）

class SimpleModule(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.a_variable = tf.Variable(5.0, name="train_me")
        self.non_trainable_variable = tf.Variable(5.0, trainable=False, name="do_not_train_me")
        
    def __call__(self, x):
        return self.a_variable * x + self.non_trainable_variable
    

关键点说明：

• class SimpleModule(tf.Module)

  ：通过 Python 的类继承机制，SimpleModule 获得了

  tf.Module

  提供的所有功能，如变量追踪与模块管理；

• super().__init__(name=name)

  ：必须显式调用父类构造函数，否则

  tf.Module

  的变量收集等功能将无法正常工作；

• tf.Variable

  的

  trainable

  参数用于标识变量是否参与训练。设置为 False 时（如固定偏置），在反向传播中不会计算梯度，常用于模型微调阶段冻结部分层；

• __call__

  方法是 Python 的魔术方法，定义后允许实例像函数一样被调用（如

  simple_module

  直接执行

  simple_module(x)

  ），而无需显式调用

  simple_module.call(x)

  ，这是深度学习框架中标准的前向推理写法。

2. 实例创建与调用

simple_module = SimpleModule(name="simple")
result = simple_module(tf.constant(5.0))
# 输出：tf.Tensor(30.0, ...) → 计算过程：5.0 * 5.0 + 5.0 = 30.0
    

上述调用实际上触发了

__call__

x

3. 查看模块变量（体现自动收集能力）

print("trainable variables:", simple_module.trainable_variables)
# 输出：(<tf.Variable 'train_me:0' ... numpy=5.0>,) → 仅包含 a_variable

print("all variables:", simple_module.variables)
    

这一部分展示了

tf.Module

tf.Module

核心价值：

tf.Module

系统能够自动汇聚所有分配给

self

的

tf.Variable

，并依据“可训练性”进行分类管理。在后续训练过程中，可直接调用

trainable_variables

进行梯度计算，无需手动遍历各个变量，极大提升了操作效率与代码简洁性。

三、第二个示例：搭建双层全连接神经网络（模块化组合结构）

本示例重点展示

tf.Module

的关键能力之一——支持子模块的递归式变量收集。当一个父模块包含多个子模块（例如模型中嵌套多个网络层）时，父模块会自动整合所有子模块中的变量，实现集中化管理。

1. 全连接层（Dense Layer）定义

全连接层（又称密集层）是构建神经网络的基本单元，其前向传播公式为：

y = ReLU(x × w + b)

其中，

w

表示权重矩阵，

b

为偏置向量，二者均为可训练参数。

class Dense(tf.Module):
    def __init__(self, in_features, out_features, name=None):
        super().__init__(name=name)
        # 权重 w：维度 [输入特征数, 输出特征数]，采用标准正态分布初始化
        self.w = tf.Variable(
            tf.random.normal([in_features, out_features]), name='w')
        # 偏置 b：维度 [输出特征数]，使用零值初始化（常见做法）
        self.b = tf.Variable(tf.zeros([out_features]), name='b')

    def __call__(self, x):
        # 执行线性变换：x 与权重 w 进行矩阵乘法后加上偏置 b
        y = tf.matmul(x, self.w) + self.b
        # 应用 ReLU 激活函数以引入非线性表达能力
        return tf.nn.relu(y)

API 说明：

• tf.random.normal([in_features, out_features])

  ：生成指定形状

[in_features, out_features]

• 的服从正态分布的随机张量，常用于权重初始化；

• tf.zeros([out_features])

  ：创建全零张量，通常用于偏置项的初始化（将偏置初始化为0是广泛采用的策略）；

• tf.matmul(x, self.w)

  ：执行矩阵乘法运算，要求左侧矩阵的列数等于右侧矩阵的行数（在此场景中，

x

• 的特征维度为

in_features

• ，与

self.w

• 的第一维匹配）；

• tf.nn.relu(y)

  ：ReLU 激活函数，当输入大于0时输出原值

y > 0

• →

y

• ，否则输出0

y ≤ 0

• ，是一种经典的非线性激活方式。

2. 构建完整模型结构（SequentialModule）

该模型由两个 Dense 子模块串联而成，形成顺序处理流程：第一层的输出作为第二层的输入。

class SequentialModule(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        # 第一隐藏层：输入维度为3，输出维度也为3
        self.dense_1 = Dense(in_features=3, out_features=3)
        # 第二输出层：输入维度需匹配上一层输出（即3），输出维度设为2
        self.dense_2 = Dense(in_features=3, out_features=2)

    def __call__(self, x):
        x = self.dense_1(x)  # 输入数据经过第一层处理
        return self.dense_2(x)  # 处理结果传入第二层并返回最终输出

其中，

self.dense_1

和

self.dense_2

均为 Dense 类的实例，同时也继承自

tf.Module

。作为容器型父模块，

SequentialModule

具备自动识别并收集其内部所有子模块变量的能力。

3. 实例化模型并执行前向推理

# 创建模型对象
my_model = SequentialModule(name="the_model")

# 输入一个 shape=(1, 3) 的张量（单个样本，含三个特征）
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))
# 输出示例：tf.Tensor([[0. 0.]], ...) —— 结果具有随机性（因权重初始值随机）

前向传播过程如下：

输入数据

[[2.0,2.0,2.0]]

→ 经过 dense_1 层（线性变换+ReLU激活）→ 得到3维中间输出 → 输入至 dense_2 层（再次线性+ReLU）→ 输出2维结果。由于权重初始化为随机值，每次运行结果可能不同。

4. 探查模型结构与变量（验证递归收集机制）

# 查看模型所包含的所有子模块（包括 dense_1 和 dense_2）
print("Submodules:", my_model.submodules)

输出结果为：(<__main__.Dense object ...>, <__main__.Dense object ...>)

通过以下代码可查看模型中包含的所有变量（包括两个子模块各自的权重 w 和偏置 b，共计4个）：

for var in my_model.variables:
    print(var, "\n")

打印出的四个变量依次为：

• dense_1 的偏置 b，形状为 shape=(3,)；
• dense_1 的权重 w，形状为 shape=(3,3)；
• dense_2 的偏置 b，形状为 shape=(2,)；
• dense_2 的权重 w，形状为 shape=(3,2)。

这一机制正是

tf.Module

所实现的「递归变量收集」功能：父级模块不仅能管理自身的参数，还会自动遍历其所有子模块，并将子模块中的变量一并纳入统一管理。这样，用户仅需通过一个模型实例即可集中访问全部可训练参数，极大简化了模型构建与维护过程。

此外，在与 Keras 框架衔接时需注意以下关键点：

• tf.Module

  是

  tf.keras.layers.Layer

  和

  tf.keras.Model

  的基类，说明 Keras 中的层与模型在底层本质上也是基于

  tf.Module

  实现的，天然具备自动变量收集和子模块管理的能力；
• 在实际开发中，可以选择使用

  tf.Module

  进行原生建模，也可以采用 Keras 提供的高层 API 快速搭建模型结构，但应避免两者混合使用，以防出现变量收集混乱或重复定义等问题；
• 无论是原生方式还是 Keras 方式，获取模型变量的方法（如

  trainable_variables

  和

  variables

  ）保持完全一致，便于统一操作和调试。

总结来看，本部分内容的核心价值在于帮助理解 TensorFlow 模型构建的底层逻辑：

• tf.Module

  作为 TensorFlow 中模型与层的基础组件，其核心设计理念是“状态（即变量）管理”与“模块化组合”；
• 它具备几项关键特性：
  • 支持按是否可训练对变量进行分类管理，便于优化器高效提取参数；
  • 能够递归遍历并整合所有子模块，适用于构建深层网络或多分支结构等复杂模型；
  • 通过

    __call__

    方法实现类函数式调用，使模型调用方式更符合深度学习编程习惯。
• 掌握

  tf.Module

  不仅为后续学习 Keras 层和模型打下坚实基础，还能深入理解模型保存、加载及训练流程背后的原理——这些操作本质上都是围绕

  tf.Module

  所收集的变量展开的管理与处理。

简而言之，

tf.Module

相当于 TensorFlow 为神经网络模型设计的一个「智能容器」，能够自动整理参数与嵌套结构，免去手动维护变量列表的繁琐工作，让用户可以更加专注于网络结构的设计与前向计算逻辑的实现。