【PyTorch】自定义神经网络层：通过继承 torch.nn.Module 类来实现

在 PyTorch 中，自定义层通常是指通过继承 torch.nn.Module 类来实现的自定义神经网络层。自定义层可以包含任意的数学操作、张量变换、初始化策略等，可以用于创建特殊的网络架构、实现非标准的操作或进行实验。

自定义层的基本步骤

自定义层通常需要以下几个步骤：

1. 继承 nn.Module：定义一个继承自 torch.nn.Module 的类。
2. 定义层的构造函数 __init__：在构造函数中定义层的参数（如权重、偏置等）。
3. 定义前向传播函数 forward：在 forward 方法中定义层的前向计算过程。

示例 1：自定义一个简单的全连接层

我们首先定义一个简单的自定义层，模拟一个具有固定权重和偏置的全连接层。

import torch
import torch.nn as nn

class MyLinearLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super(MyLinearLayer, self).__init__()
        # 定义权重和偏置
        self.weights = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        # 实现前向传播
        return torch.matmul(x, self.weights.T) + self.bias

# 使用自定义层
model = MyLinearLayer(3, 2)  # 输入3个特征，输出2个特征
input_data = torch.randn(4, 3)  # 4个样本，每个样本3个特征
output_data = model(input_data)

print(output_data)

解释

• __init__：在 __init__ 中，我们定义了两个参数：weights（权重矩阵）和 bias（偏置向量），并将它们转换为 nn.Parameter 类型，使它们成为可训练的参数。
• forward：在 forward 方法中，我们手动实现了线性变换：y = x * W^T + b，其中 x 是输入张量，W 是权重，b 是偏置。

示例 2：自定义一个激活函数

下面我们定义一个自定义的激活函数 MyReLU，它实现了一个简单的 ReLU（Rectified Linear Unit）激活操作。

class MyReLU(nn.Module):
    def __init__(self):
        super(MyReLU, self).__init__()

    def forward(self, x):
        # 实现ReLU激活
        return torch.maximum(x, torch.zeros_like(x))

# 使用自定义激活函数
model = nn.Sequential(
    nn.Linear(3, 2),  # 全连接层
    MyReLU()          # 自定义ReLU层
)

input_data = torch.randn(4, 3)  # 4个样本，每个样本3个特征
output_data = model(input_data)

print(output_data)

解释

• MyReLU 层的 forward 方法实现了 ReLU 激活操作。我们使用 torch.maximum(x, torch.zeros_like(x)) 来保证所有负值都被置为零。

示例 3：自定义一个卷积层

下面是一个自定义卷积层的例子，它实现了一个简单的 2D 卷积操作。我们通过直接操作卷积核（filters）和步幅（stride）来实现卷积。

class MyConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(MyConv2d, self).__init__()
        # 定义卷积核参数
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))
        self.bias = nn.Parameter(torch.zeros(out_channels))

    def forward(self, x):
        # 简单的卷积操作 (不考虑 padding 和 stride)
        batch_size, in_channels, height, width = x.shape
        out_channels, _, kernel_height, kernel_width = self.weight.shape
        
        # 输出尺寸计算
        out_height = height - kernel_height + 1
        out_width = width - kernel_width + 1
        
        # 创建输出张量
        output = torch.zeros(batch_size, out_channels, out_height, out_width)
        
        # 执行卷积操作
        for i in range(out_height):
            for j in range(out_width):
                region = x[:, :, i:i+kernel_height, j:j+kernel_width]
                output[:, :, i, j] = torch.sum(region * self.weight, dim=(1, 2, 3)) + self.bias
                
        return output

# 使用自定义卷积层
model = MyConv2d(1, 2, 3)  # 输入1个通道，输出2个通道，卷积核大小为 3x3
input_data = torch.randn(4, 1, 5, 5)  # 4个样本，每个样本 1个通道，5x5的图像
output_data = model(input_data)

print(output_data.shape)  # 输出尺寸 (4, 2, 3, 3)

解释

• MyConv2d 类实现了一个简单的卷积层。我们手动计算卷积输出，通过滑动窗口（对于每个输出位置，取对应的输入区域与卷积核的元素逐点相乘并求和）。
• 在实际应用中，PyTorch 提供了高效的卷积操作（如 nn.Conv2d），该例子只是为了展示如何手动实现卷积操作。

示例 4：自定义层中使用已有层

自定义层还可以在内部使用已有的 PyTorch 层。例如，我们可以自定义一个包含多个层的复合层，如卷积层后跟 ReLU 激活。

class ConvReLUBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ConvReLUBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

# 使用自定义的 ConvReLUBlock
model = nn.Sequential(
    ConvReLUBlock(1, 16, 3),  # 输入1个通道，输出16个通道，卷积核大小为3x3
    nn.MaxPool2d(2)           # 池化层
)

input_data = torch.randn(4, 1, 5, 5)  # 4个样本，1个通道，5x5的图像
output_data = model(input_data)

print(output_data.shape)  # 输出尺寸 (4, 16, 1, 1)

解释

• ConvReLUBlock 类定义了一个包含卷积层和 ReLU 激活函数的复合层。我们可以将多个已有的 PyTorch 层组合到一个自定义层中。
• 使用 nn.Sequential 方便地将多个层组合成一个模块。

注意事项

• 参数初始化：当你定义了自己的层时，通常需要考虑如何初始化层的参数。可以通过 torch.nn.init 中的各种方法（如 xavier_uniform_、kaiming_normal_ 等）来初始化参数。

• 可训练参数：自定义层的参数需要通过 nn.Parameter 来声明，这样 PyTorch 才能追踪它们并在训练过程中进行更新。

• 前向传播：forward 方法定义了层的前向计算流程，可以包含任意的张量操作（如加法、乘法、非线性变换等）。

总结

• 自定义层：通过继承 torch.nn.Module，我们可以创建复杂的自定义神经网络层。
• 可训练参数：自定义层中的权重和偏置等可训练参数需要通过 nn.Parameter 来声明。
• 前向传播：在 forward 方法中，我们定义了层的计算过程，PyTorch 会自动处理反向传播和梯度更新。