Leaky ReLU 激活函数的 PyTorch 自定义参数实现

Leaky ReLU 的数学表达式为： $$f(x) = \begin{cases} x & \text{if } x \geq 0 \ \alpha x & \text{if } x < 0 \end{cases}$$ 其中 $\alpha$ 是负半轴的斜率参数（通常取 0.01）。

自定义实现步骤

1. 导入 PyTorch 模块

import torch
import torch.nn as nn

1. 创建可学习参数类

class LearnableLeakyReLU(nn.Module):
    def __init__(self, alpha_init=0.01):
        super().__init__()
        # 初始化可学习参数，设置为可训练
        self.alpha = nn.Parameter(torch.tensor(alpha_init))
    
    def forward(self, x):
        # 使用 torch.where 实现条件计算
        return torch.where(x >= 0, x, self.alpha * x)

1. 梯度验证测试

def test_gradient():
    # 创建输入张量（包含正负数）
    x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0], requires_grad=True)
    
    # 实例化激活函数（初始α=0.1）
    leaky_relu = LearnableLeakyReLU(alpha_init=0.1)
    
    # 前向传播
    y = leaky_relu(x)
    print("前向输出:", y.detach().numpy())  # [-0.2, -0.1, 0.0, 1.0, 2.0]
    
    # 反向传播
    y.sum().backward()
    print("输入梯度:", x.grad.numpy())      # [0.1, 0.1, 1.0, 1.0, 1.0]
    print("α参数梯度:", leaky_relu.alpha.grad.item())  # -3.0 (∑x_neg)

1. 网络集成示例

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            LearnableLeakyReLU(0.01),  # 使用自定义激活函数
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3),
            LearnableLeakyReLU(0.01)   # 第二层激活
        )
        self.fc = nn.Linear(32*6*6, 10)
    
    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

关键特性说明

1. 参数学习机制：

   • $\alpha$ 通过梯度下降自动优化
   • 反向传播时梯度计算： $$\frac{\partial L}{\partial \alpha} = \sum_{x_i<0} \frac{\partial L}{\partial y_i} \cdot x_i$$

2. 训练建议：

   model = CNN()
   optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
   
   # 训练循环中自动更新α参数
   for epoch in range(10):
       for data, target in dataloader:
           optimizer.zero_grad()
           output = model(data)
           loss = nn.CrossEntropyLoss()(output, target)
           loss.backward()
           optimizer.step()
   

3. 参数初始化技巧：

   • 推荐初始值：$\alpha_{\text{init}} \in [0.01, 0.3]$
   • 避免初始值 $\alpha_{\text{init}} = 0$（梯度消失）
   • 可通过 nn.init.constant_(module.alpha, 0.1) 显式初始化

此实现允许模型自动学习最优的负半轴斜率，相比固定参数的 ReLU 变体能更好地适应复杂数据分布。实际应用中可观察到 $\alpha$ 值随训练过程动态调整，通常在 0.01~0.2 区间收敛。