在 PyTorch 中构建神经网络模型涉及理解 torch.nn.Module，它作为所有组件的基础。本实践指南专注于将 Keras 模型架构直接转换为其 PyTorch 等效形式，并与 Keras 层和模型构建 API 进行比较。通过完成这些示例，你将巩固对如何在 PyTorch 的 nn.Module 框架内定义层、构建模型和管理数据流的理解。

我们将从简单的模型开始，逐步过渡到能够体现 PyTorch “按运行定义” 方法灵活性的结构，这种灵活性在构建非严格序列模型时尤为明显。

首先，请确保已导入 PyTorch：

import torch
import torch.nn as nn
import torch.nn.functional as F # 经常用于激活函数

示例 1：简单前馈网络（全连接层）

让我们从一个基本的全连接神经网络开始，这是许多教程中的常见起点。

Keras 序列模型：

在 Keras 中，你可能会这样定义一个简单的两层网络，例如用于 MNIST 数字分类（假设输入已展平）：

# TensorFlow/Keras
# import tensorflow as tf
#
# model_keras_ffn = tf.keras.Sequential([
#     tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
#     tf.keras.layers.Dropout(0.2),
#     tf.keras.layers.Dense(10) # 输出层，激活函数通常由损失函数处理
# ])
#
# model_keras_ffn.summary()

input_shape 在第一层中定义。最终的激活函数（例如，用于分类的 softmax）有时会被省略，如果损失函数（如 tf.keras.losses.CategoricalCrossentropy(from_logits=True)）期望原始对数（logits）。

PyTorch nn.Module 等效形式：

要在 PyTorch 中构建它，我们将继承 nn.Module：

class PyTorchSimpleFFN(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes, dropout_rate=0.2):
        super(PyTorchSimpleFFN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout_rate)
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # 期望 x 的形状为 (batch_size, input_size)
        out = self.fc1(x)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc2(out)
        # 注意：Softmax 通常在后续应用，例如在 nn.CrossEntropyLoss 内部
        return out

# 实例化模型
input_size = 784
hidden_size = 128
num_classes = 10
pytorch_model_ffn = PyTorchSimpleFFN(input_size, hidden_size, num_classes)
print(pytorch_model_ffn)

# 带有模拟数据的示例用法：
dummy_input_ffn = torch.randn(64, input_size) # 批次大小为 64，784 个特征
output_ffn = pytorch_model_ffn(dummy_input_ffn)
print("Output shape:", output_ffn.shape) # 期望：torch.Size([64, 10])

说明：

__init__ 中的层：PyTorch 层（例如 nn.Linear、nn.ReLU、nn.Dropout）通常在 __init__ 方法中定义为属性。nn.Linear 等效于 Keras 的 Dense 层。
明确的 forward 方法：forward 方法明确定义了输入数据如何流经各层。这体现了 PyTorch 动态特性的优点。
输入形状：与 Keras 可以在第一层指定 input_shape 不同，PyTorch 模型通常会适应其在 forward 方法中接收到的输入形状。nn.Linear 层的 in_features 必须与输入张量的特征维度匹配。
激活函数：nn.ReLU 是一个 nn.Module，因此它在 __init__ 中实例化。另外，F.relu（来自 torch.nn.functional）可以直接在 forward 中应用，无需预先实例化。
输出层激活：与 Keras 类似，如果你在 PyTorch 中使用 nn.CrossEntropyLoss 等损失函数（它结合了 nn.LogSoftmax 和 nn.NLLLoss），通常无需在模型内部对输出层应用 softmax 激活。原始分数（logits）会传递给损失函数。

示例 2：基本卷积神经网络 (CNN)

CNN 对图像处理非常重要。让我们转换一个简单的 Keras CNN。

Keras 序列 CNN：

# TensorFlow/Keras
# import tensorflow as tf
#
# model_keras_cnn = tf.keras.Sequential([
#     tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(28, 28, 1), padding='same'),
#     tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
#     tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation='relu', padding='same'),
#     tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
#     tf.keras.layers.Flatten(),
#     tf.keras.layers.Dense(128, activation='relu'),
#     tf.keras.layers.Dense(10) # 输出 logits
# ])
#
# model_keras_cnn.summary()

这个 Keras 模型假定输入形状为 (height, width, channels)。

PyTorch nn.Module 等效形式：

PyTorch 期望图像数据格式为 (batch_size, channels, height, width)。

class PyTorchSimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(PyTorchSimpleCNN, self).__init__()
        # 卷积层 1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1) # padding='same'
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 卷积层 2
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1) # padding='same'
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 展平层和全连接层
        # 为了确定 fc1 的输入大小，我们需要计算卷积和池化层之后的输出形状。
        # 假设输入为 28x28：
        # 经过 conv1 (32, 28, 28)，pool1 (32, 14, 14)
        # 经过 conv2 (64, 14, 14)，pool2 (64, 7, 7)
        # 展平后大小：64 * 7 * 7
        self.fc1_input_features = 64 * 7 * 7
        self.fc1 = nn.Linear(self.fc1_input_features, 128)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        # x 形状：(batch_size, 1, 28, 28)
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(self.conv2(x)))

        # 为全连接层展平输出
        # x.size(0) 是批次大小。-1 表示推断其余维度。
        x = x.view(x.size(0), -1) # 或者使用 torch.flatten(x, start_dim=1)

        x = self.relu3(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化模型
pytorch_model_cnn = PyTorchSimpleCNN(num_classes=10)
print(pytorch_model_cnn)

# 带有模拟数据的示例用法：
dummy_input_cnn = torch.randn(64, 1, 28, 28) # 批次大小为 64，1 个通道，28x28 图像
output_cnn = pytorch_model_cnn(dummy_input_cnn)
print("Output shape:", output_cnn.shape) # 期望：torch.Size([64, 10])

说明：

通道顺序：请记住 PyTorch 图像张量的 N C H W（批次、通道、高、宽）约定，与 TensorFlow 的默认 N H W C 不同。
nn.Conv2d：参数包括 in_channels、out_channels、kernel_size、stride、padding。当步长为 1 时，kernel_size=3 且 padding=1 通常与 Keras 的 padding='same' 效果相近。
展平：从卷积/池化层到全连接层的转换需要展平。在 PyTorch 中，x.view(x.size(0), -1) 或 torch.flatten(x, start_dim=1) 是常见做法。必须计算最后一个池化层之后的特征数量，以便为第一个 nn.Linear 层正确设定大小。这是 PyTorch 中的一个手动步骤，除非使用自适应池化层（例如 nn.AdaptiveAvgPool2d((1,1))），它们无论输入大小如何都能输出固定大小的特征图，从而简化了后续的展平层和线性层定义。
nn.MaxPool2d：kernel_size 和 stride 是重要的参数。

示例 3：带有跳跃连接的模型（函数式 API 风格）

Keras 的函数式 API 允许构建更复杂的架构，例如具有多输入、多输出或跳跃连接的架构。PyTorch 的 nn.Module 通过 forward 方法天然支持这种灵活性。让我们构建一个带有残差（跳跃）连接的简单块。

Keras 函数式 API 模型（示例块）：

# TensorFlow/Keras
# import tensorflow as tf
#
# input_tensor = tf.keras.Input(shape=(64, 64, 3))
# x = tf.keras.layers.Conv2D(32, (3,3), padding='same', activation='relu')(input_tensor)
# x = tf.keras.layers.Conv2D(32, (3,3), padding='same')(x) # 暂无激活函数
#
# # 示例：简化的残差连接
# # 对于真正的 ResNet 块，通道维度可能需要匹配（例如，使用 1x1 卷积）
# # 这里，我们假设 input_tensor 和 x 在卷积后具有兼容的形状以进行加法运算
# # 或者，如果通道不同，则投影 input_tensor
# identity = tf.keras.layers.Conv2D(32, (1,1), padding='same')(input_tensor) # 投影恒等映射
#
# added = tf.keras.layers.Add()([x, identity])
# output_tensor = tf.keras.layers.Activation('relu')(added)
#
# model_keras_functional = tf.keras.Model(inputs=input_tensor, outputs=output_tensor)
# model_keras_functional.summary()

PyTorch nn.Module 等效形式（简化残差块）：

我们将创建一个模块来实现在卷积层输出中加入输入的简化残差块。为简单起见，我们假设通道数量保持不变，或者，如果维度需要匹配，我们将使用 1x1 卷积进行恒等映射路径。

class PyTorchResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(PyTorchResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels) # 批归一化在 ResNet 中很常见
        self.relu = nn.ReLU(inplace=True) # inplace=True 可以节省内存
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)

        # 快捷连接（恒等映射或投影）
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        identity = self.shortcut(x) # 将快捷变换应用于 x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity # 逐元素相加
        out = self.relu(out)
        return out

# 实例化块
# 示例：输入 3 通道，输出 64 通道，下采样
res_block = PyTorchResidualBlock(in_channels=3, out_channels=64, stride=2)
print(res_block)

# 带有模拟数据的示例用法：
dummy_input_res = torch.randn(16, 3, 224, 224) # 批次大小为 16，3 个通道，224x224 图像
output_res = res_block(dummy_input_res)
# 如果 stride=2，H 和 W 将减半。输出通道将是 64。
print("Output shape:", output_res.shape) # 期望：torch.Size([16, 64, 112, 112])

# 示例：通道或维度无变化
res_block_same_dim = PyTorchResidualBlock(in_channels=64, out_channels=64, stride=1)
dummy_input_same_dim = torch.randn(16, 64, 56, 56)
output_same_dim = res_block_same_dim(dummy_input_same_dim)
print("Output shape (same dim):", output_same_dim.shape) # 期望：torch.Size([16, 64, 56, 56])

说明：

任意 forward 逻辑：forward 方法可以实现任何计算，包括跳跃连接、分支或自定义操作。这充分体现了 PyTorch 的灵活性。你直接定义数据流向。
nn.Sequential 用于快捷连接：nn.Sequential 可用于组合层，例如，在快捷连接中，如果需要投影（如 1x1 卷积）来匹配维度。
参数共享：如果你想多次重用同一层的同一个实例（共享权重），你会在 __init__ 中定义一次，并在 forward 中多次调用。如果你需要同类型层的不同实例，则在 __init__ 中将它们定义为单独的属性。

验证等效性

尽管我们构建了“等效”模型，真正的等效意味着在相同权重下，它们对相同输入产生相同的输出。

架构：主要检查是层序列、它们的类型以及它们的参数（核大小、单元/滤波器数量、步长、填充）是否匹配。在 PyTorch 和 Keras 中打印模型结构（使用 model.summary()）有助于确认。
参数数量：可训练参数的数量应非常接近。细微的差异可能源于偏置的默认处理方式或特定的层实现，但主要的差异表明架构不匹配。
权重初始化：TensorFlow/Keras 和 PyTorch 之间的默认权重初始化方案可能不同。为了获得相同的输出（用于测试目的），你需要手动将所有权重设置为相同。这对于一般的模型转换通常不是必需的，因为模型反正都是从随机初始化开始训练的。
前向传播测试：输入一个通用模拟数据（针对每个框架进行适当格式化）并检查输出形状是一个很好的合理性验证。

这些实践示例展示了将 Keras 模型转换为 PyTorch 的核心过程。核心在于理解如何在 __init__ 中定义层，然后如何在 forward 方法中明确组织它们的执行。如你所见，PyTorch 的 nn.Module 提供了一种非常灵活且符合 Python 习惯的方式来定义复杂的模型架构。下一步是学习如何将数据输入这些模型并进行训练。

这部分内容有帮助吗？