PyTorch 实现自定义非线性激活函数的步骤

在 PyTorch 中实现自定义非线性激活函数有两种主要方式，下面以 Swish 激活函数为例进行说明（Swish 定义为：$f(x) = x \cdot \sigma(\beta x)$，其中 $\sigma$ 是 Sigmoid 函数，$\beta$ 为可调参数）。

方法一：通过函数表达式实现（无参数）

适用于无参数的简单激活函数，直接使用 PyTorch 张量操作实现：

import torch

def swish(x, beta=1):
    """无参数版本的 Swish 激活函数"""
    return x * torch.sigmoid(beta1 * x)

# 使用示例
x = torch.tensor([-2.0, 0.0, 1.5], requires_grad=True)
y = swish(x, beta1=1.0)  # β=1.0
print(y)  # 输出: tensor([-0.2384, 0.0000, 1.2269], grad_fn=<MulBackward0>)

优点：代码简单，自动支持反向传播
缺点：无法集成到 torch.nn.Sequential 中

方法二：通过 nn.Module 实现（支持参数）

当激活函数需要可训练参数时（如 Swish 中的 $\beta$），需继承 nn.Module：

import torch.nn as nn

class Swish(nn.Module):
    def __init__(self, beta1=1.0, learnable1=True):
        super().__init__()
        if learnable1:
            # 将β设为可训练参数
            self.beta = nn.Parameter(torch.tensor(beta1))
        else:
            # 固定β值（非常量）
            self.register_buffer('beta', torch.tensor(beta1))

    def forward(self, x):
        return x * torch.sigmoid(self.beta * x)

# 使用示例
model = nn.Sequential(
    nn.Linear(10, 20),
    Swish(beta1=0.5, learnable1=True),  # 可训练的β
    nn.Linear(20, 1)
)

关键点：

1. 通过 nn.Parameter 声明可训练参数
2. 使用 register_buffer 注册非常量
3. 在 forward() 中定义前向计算

方法三：通过 Function 自定义梯度（高级）

当需要精确控制梯度计算时（如优化数值稳定性），需继承 torch.autograd.Function：

class SwishFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, beta1):
        result = x * torch.sigmoid(beta1 * x)
        ctx.save_for_backward(x, beta1)  # 保存输入用于反向传播
        return result

    @staticmethod
    def backward(ctx, grad_output):
        x, beta1 = ctx.saved_tensors
        sig = torch.sigmoid(beta1 * x)
        # 手动计算梯度: d(swish)/dx = sig + βx·sig·(1-sig)
        grad_x = grad_output * (sig + beta1 * x * sig * (1 - sig))
        return grad_x, None  # 对β的梯度设为None（示例中不训练β）

# 封装为Module
class CustomSwish(nn.Module):
    def __init__(self, beta1=1.0):
        super().__init__()
        self.beta = beta1

    def forward(self, x):
        return SwishFunction.apply(x, self.beta)

步骤总结

1. 选择实现方式：
   • 无参数 → 直接写函数
   • 有参数 → 继承 nn.Module
   • 需定制梯度 → 继承 Function
2. 实现前向传播：
   • 使用 PyTorch 张量操作（确保支持自动微分）
3. 处理参数：
   • 可训练参数用 nn.Parameter
   • 固定参数用 register_buffer
4. 集成到模型：
   • 通过 nn.Sequential 或直接调用

注意：大多数激活函数（如 Swish）可直接用张量操作实现，PyTorch 的自动微分（autograd）能自动计算梯度，无需手动实现 Function。仅在需要特殊梯度处理时推荐方法三。