python - Pytorch 自定义激活函数？

Question

我在 Pytorch 中实现自定义激活函数时遇到问题，例如 Swish。我应该如何在 Pytorch 中实现和使用自定义激活函数？

Answer 1

有四种可能性，具体取决于您要查找的内容。您需要问自己两个问题:

Q1)你的激活函数会有可学习的参数吗？

如果是，您别无选择，只能将激活函数创建为 nn.Module 类，因为您需要存储这些权重。

如果否，您可以自由地创建一个普通函数或一个类，具体取决于您方便什么。

Q2)您的激活函数能否表示为现有 PyTorch 函数的组合？

如果是，您可以简单地将其编写为现有 PyTorch 函数的组合，而无需创建定义梯度的 backward 函数。

如果否，您将需要手写渐变。

例1:SiLU函数

SiLU function f(x) = x * sigmoid(x) 没有任何学习权重，可以完全用现有的 PyTorch 函数编写，因此您可以简单地将其定义为函数:

def silu(x):
    return x * torch.sigmoid(x)

然后像使用 torch.relu 或任何其他激活函数一样简单地使用它。

示例 2:具有学习斜率的 SiLU

在这种情况下，您有一个学习参数，即斜率，因此您需要对其进行分类。

class LearnedSiLU(nn.Module):
    def __init__(self, slope = 1):
        super().__init__()
        self.slope = slope * torch.nn.Parameter(torch.ones(1))

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

示例3:带向后

如果您需要创建自己的渐变函数，可以查看此示例:Pytorch: define custom function