定义网络组件：Keras 层与 torch.nn.Module

从 TensorFlow Keras 转向 PyTorch 时，了解网络组件的定义方式是首要的一步。tf.keras.layers.Layer 是 Keras 中所有层的基类，这些层是模型的基本构成单元。PyTorch 通过其 torch.nn.Module 类提供了一个相似但有所不同的基准。每个神经网络模块，从单个层到整个复杂模型，都是 nn.Module 的子类。该类提供了构建神经网络所需的核心功能：它可以注册和管理参数（如权重和偏置），容纳子模块（其他 nn.Module 实例），并定义前向计算。

torch.nn.Module 类：你的网络蓝图

将 torch.nn.Module 视为神经网络中任何具有可学习参数或执行独立计算步骤的部分的蓝图。其主要职责包括：

1. 参数管理：nn.Module 可以包含 torch.nn.Parameter 实例。这些是特殊的张量，它们会被自动注册为模型参数。当你要求模型提供其参数时（例如，传递给优化器），nn.Module 会确保所有这些已注册的参数，包括嵌套子模块中的参数，都是可访问的。
2. 子模块组织：一个 nn.Module 可以将其他 nn.Module 实例作为属性包含。这使你能够通过分层方式组合更简单的模块来构建复杂的架构。
3. 定义计算：每个 nn.Module 子类通常会实现一个 forward() 方法。此方法接收输入张量并执行已定义的操作，返回输出张量。这与 Keras Layer 中的 call() 方法直接对应。

尽管 Keras 层有一个 build() 方法，常用于延迟权重创建（即在首次得知输入形状时创建权重），但 PyTorch 模块通常在其 __init__() 构造函数中直接定义其层以及参数的形状。像 nn.Linear 这样的预构建层的权重在实例化时创建，这要求你届时指定输入特征维度。

在 __init__ 中定义层和参数

在 PyTorch 中，标准做法是在模块的 __init__ 构造函数中定义并初始化所有组成层和参数。当你将另一个 nn.Module 的实例（如 nn.Linear 或 nn.Conv2d）作为自定义模块的属性时，PyTorch 会自动将其识别为子模块。

让我们看一个简单例子。如果你想创建一个包含线性层和 ReLU 激活函数的模块：

import torch
import torch.nn as nn

class SimplePyTorchModule(nn.Module):
    def __init__(self, input_features, output_features):
        super(SimplePyTorchModule, self).__init__()
        # 在此处定义子模块（层）
        self.linear_layer = nn.Linear(input_features, output_features)
        self.activation = nn.ReLU()

        # 自定义可学习参数示例
        # self.my_custom_bias = nn.Parameter(torch.zeros(output_features))

    def forward(self, x):
        # 使用子模块定义计算流程
        x = self.linear_layer(x)
        x = self.activation(x)
        # 如果 self.my_custom_bias 已定义：
        # x = x + self.my_custom_bias 
        return x

# 实例化模块
module = SimplePyTorchModule(input_features=10, output_features=5)
print(module)

输出:

SimplePyTorchModule(
  (linear_layer): Linear(in_features=10, out_features=5, bias=True)
  (activation): ReLU()
)

在这个 __init__ 方法中：

• super(SimplePyTorchModule, self).__init__() 对于调用基类 nn.Module 的构造函数非常重要。
• self.linear_layer = nn.Linear(input_features, output_features) 创建了 PyTorch 预构建线性层的一个实例，并将其作为属性分配。它的参数（权重和偏置）会自动注册到 SimplePyTorchModule。
• self.activation = nn.ReLU() 创建了 ReLU 激活函数的一个实例。ReLU 本身没有可学习参数，但它仍然是一个 nn.Module。

这类似于你在自定义 tf.keras.Layer 子类的 __init__ 或 build 方法中定义层的方式。然而，组合成更大的模型通常有所不同。在 Keras 中，你可能会将这些添加到 tf.keras.Sequential 模型中，或使用函数式 API 连接它们。在 PyTorch 中，你通常通过嵌套直接构建这些 nn.Module 结构。

实现 forward() 方法

forward() 方法是你定义模块实际计算的地方。它接受一个或多个输入张量，并返回一个或多个输出张量。你使用在 __init__ 中定义的层和参数来转换输入。

继续 SimplePyTorchModule：

# （在 SimplePyTorchModule 类内部）
# def forward(self, x):
#     x = self.linear_layer(x)
#     x = self.activation(x)
#     return x

这里，self.linear_layer(x) 调用了 nn.Linear 实例的 forward 方法。nn.Module 实例是可调用的；当你调用 module_instance(input) 时，它会在内部调用 module_instance.forward(input)。

PyTorch 的一个重要特点是其动态计算图（即“按运行定义”）。forward() 方法只是普通的 Python 代码。这意味着你可以使用标准的 Python 控制流（循环、条件语句）来定义复杂、自适应的计算。运算图是在 forward() 方法执行时即时构建的。这与 TensorFlow 的传统图模式（即“按定义运行”）形成对比，在传统模式下，图通常是先静态定义的（尽管 TensorFlow 2.x 默认的 Eager Execution 模式行为更像 PyTorch 的“按运行定义”）。

可学习参数：torch.nn.Parameter

尽管像 nn.Linear 这样的预构建层管理着自己的参数，但有时你需要直接在模块中定义自己的自定义可学习参数。为此，你需要用 torch.nn.Parameter 封装一个 torch.Tensor。这会告诉 PyTorch 该张量应被视为模型参数，意味着它的 requires_grad 属性默认将被设置为 True（因此在反向传播期间会为其计算梯度），并且它将被包含在 model.parameters() 返回的列表中。

import torch
import torch.nn as nn

class ModuleWithCustomParameter(nn.Module):
    def __init__(self, num_features):
        super(ModuleWithCustomParameter, self).__init__()
        # 每个特征的可学习缩放因子
        self.scale = nn.Parameter(torch.ones(num_features)) 
        # 每个特征的可学习偏置
        self.bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x):
        # x 的预期形状为 (batch_size, num_features)
        return x * self.scale + self.bias

custom_param_module = ModuleWithCustomParameter(5)
# 你可以查看其参数：
for name, param in custom_param_module.named_parameters():
    print(f"{name}: {param.data}") 

这与在 Keras Layer 的 build 方法中使用 self.add_weight() 来创建和注册可训练权重相似。

嵌套模块以构建分层架构

nn.Module 的强大能力之一是它能够包含其他 nn.Module 实例。这使你能够通过分层方式组合更简单、可重用的块来构建复杂模型。父模块会自动发现并管理其嵌套子模块中的所有参数。

考虑构建一个稍微复杂一些的网络，它使用我们之前定义的 SimplePyTorchModule 作为构建块：

class AdvancedNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim1, hidden_dim2, output_dim):
        super(AdvancedNetwork, self).__init__()
        # 将 SimplePyTorchModule 嵌套为一个“块”
        self.block1 = SimplePyTorchModule(input_dim, hidden_dim1)
        self.intermediate_linear = nn.Linear(hidden_dim1, hidden_dim2)
        self.relu = nn.ReLU()
        self.output_layer = nn.Linear(hidden_dim2, output_dim)

    def forward(self, x):
        x = self.block1(x) # 使用 SimplePyTorchModule 的前向传播
        x = self.intermediate_linear(x)
        x = self.relu(x)
        x = self.output_layer(x)
        return x

# 实例化示例
adv_net = AdvancedNetwork(input_dim=20, hidden_dim1=15, hidden_dim2=10, output_dim=2)
print(adv_net)

输出:

AdvancedNetwork(
  (block1): SimplePyTorchModule(
    (linear_layer): Linear(in_features=20, out_features=15, bias=True)
    (activation): ReLU()
  )
  (intermediate_linear): Linear(in_features=15, out_features=10, bias=True)
  (relu): ReLU()
  (output_layer): Linear(in_features=10, out_features=2, bias=True)
)

如你所见，AdvancedNetwork 包含了 block1，它是 SimplePyTorchModule 的一个实例。来自 block1.linear_layer、intermediate_linear 和 output_layer 的所有参数都将成为 adv_net.parameters() 的一部分。这种可组合性是有效组织 PyTorch 代码的核心所在。

下图说明了一个自定义 PyTorch nn.Module 的结构，该模块包含其他模块（包括预构建的和其他可能的自定义模块）和参数。

一个 PyTorch nn.Module (MyNetwork) 在 __init__ 中定义了其组件（子模块如 CustomBlockA、nn.ReLU、nn.Linear，以及直接的 nn.Parameter 如 global_bias），并在 forward 中定义了它们的计算流程。CustomBlockA 本身是另一个 nn.Module，展示了嵌套的用法。

torch.nn 中的预构建层

PyTorch 在 torch.nn 包中提供了一系列丰富的预构建层，例如：

• nn.Linear：全连接层。
• nn.Conv1d、nn.Conv2d、nn.Conv3d：用于不同维度的卷积层。
• nn.RNN、nn.LSTM、nn.GRU：循环层。
• nn.BatchNorm1d、nn.BatchNorm2d：批归一化层。
• nn.Dropout：Dropout 层。
• 激活函数，如 nn.ReLU、nn.Sigmoid、nn.Tanh、nn.Softmax（尽管许多激活函数在 torch.nn.functional 中也可用，并且可以直接在 forward 方法中应用）。

这些预构建层本身都是 nn.Module 的子类。你通过在模块的 __init__ 方法中实例化它们，然后在 forward 方法中使用适当的输入来调用它们。这与你在 TensorFlow 中使用 tf.keras.layers 中的层的方式相似。后续章节将更详细地介绍这些常见层类型，比较它们的 PyTorch 实现与 Keras 对应部分。

从 Keras 的 Layer 转向 PyTorch 的 nn.Module 需要接受一种结构，在这种结构中，你在 __init__ 中明确定义网络模块的组成部分，并在 forward 中定义它们的计算流程。这种明确的定义提供了细粒度控制，并与 Python 的动态能力自然结合，使得构建和调试即使非常复杂的模型架构也变得简单。