如何优化神经网络激活函数及其导数获取代码并提升Python风格？

一、核心优化方向 #

原代码的主要冗余点在于每个match-case分支重复定义函数并返回，可用函数列表available_fn与分支名称存在重复维护的问题，容易出现不一致。我们可以通过字典映射统一管理激活函数对的方式，让代码更简洁、易维护，更符合Python风格。

二、优化后的字典映射版本 #

这种方式把所有激活函数及其导数提前定义，用字典建立名称与函数对的映射，逻辑更清晰，避免重复代码：

import numpy as np

def identity(x):
    return x

def identity_deriv(x):
    return np.ones(x.shape)

def relu(x):
    return np.maximum(x, 0.0)

def relu_deriv(x):
    return np.where(x >= 0.0, 1.0, 0.0)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def sigmoid_deriv(x):
    sig_x = sigmoid(x)
    return sig_x * (1 - sig_x)

def tanh(x):
    exp_2x = np.exp(2 * x)
    return (exp_2x - 1) / (exp_2x + 1)

def tanh_deriv(x):
    return 1 - tanh(x) ** 2

_ACTIVATION_MAP = {
    'identity': (identity, identity_deriv),
    'relu': (relu, relu_deriv),
    'sigmoid': (sigmoid, sigmoid_deriv),
    'tanh': (tanh, tanh_deriv)
}

def get_activation_fn_with_deriv(fn_name):
    """返回指定激活函数及其导数的函数对象

    Args:
        fn_name: 激活函数名称
    
    Returns:
        (fn, fn_deriv): 激活函数及其导数的元组
    
    Raises:
        ValueError: 请求的激活函数不存在时抛出
    
    Examples:
        >>> relu, relu_deriv = get_activation_fn_with_deriv('relu')
        >>> relu(np.array([1, -1, 2.3]))
        array([1. , 0. , 2.3])
    """
    fn_key = fn_name.lower()
    if fn_key not in _ACTIVATION_MAP:
        available = list(_ACTIVATION_MAP.keys())
        raise ValueError(f"指定的激活函数不可用，可选范围：{available}")
    return _ACTIVATION_MAP[fn_key]

三、类封装的可行性分析 #

将激活函数封装到类中是可行的，且在某些场景下更优：

• 优势：可以把激活函数的计算、导数、名称甚至其他属性（如是否支持批量计算、初始化建议）封装在一起，结构更规整，便于后续扩展（比如新增激活函数时，只需新增一个类）。
• 劣势：如果只是单纯需要获取函数对，类封装会比字典映射更繁琐，增加了不必要的面向对象开销。

下面是类封装的示例：

import numpy as np

class ActivationFunction:
    def __call__(self, x):
        raise NotImplementedError("需实现激活函数计算逻辑")
    
    def deriv(self, x):
        raise NotImplementedError("需实现导数计算逻辑")

class Identity(ActivationFunction):
    def __call__(self, x):
        return x
    
    def deriv(self, x):
        return np.ones(x.shape)

class ReLU(ActivationFunction):
    def __call__(self, x):
        return np.maximum(x, 0.0)
    
    def deriv(self, x):
        return np.where(x >= 0.0, 1.0, 0.0)

class Sigmoid(ActivationFunction):
    def __call__(self, x):
        return 1 / (1 + np.exp(-x))
    
    def deriv(self, x):
        sig_x = self(x)
        return sig_x * (1 - sig_x)

class Tanh(ActivationFunction):
    def __call__(self, x):
        exp_2x = np.exp(2 * x)
        return (exp_2x - 1) / (exp_2x + 1)
    
    def deriv(self, x):
        return 1 - self(x) ** 2

_ACTIVATION_CLASS_MAP = {
    'identity': Identity,
    'relu': ReLU,
    'sigmoid': Sigmoid,
    'tanh': Tanh
}

def get_activation_fn_with_deriv(fn_name):
    """返回指定激活函数及其导数的实例

    Args:
        fn_name: 激活函数名称
    
    Returns:
        激活函数实例，可通过`instance(x)`调用激活函数，`instance.deriv(x)`调用导数
    
    Raises:
        ValueError: 请求的激活函数不存在时抛出
    
    Examples:
        >>> relu = get_activation_fn_with_deriv('relu')
        >>> relu(np.array([1, -1, 2.3]))
        array([1. , 0. , 2.3])
        >>> relu.deriv(np.array([1, -1, 2.3]))
        array([1., 0., 1.])
    """
    fn_key = fn_name.lower()
    if fn_key not in _ACTIVATION_CLASS_MAP:
        available = list(_ACTIVATION_CLASS_MAP.keys())
        raise ValueError(f"指定的激活函数不可用，可选范围：{available}")
    return _ACTIVATION_CLASS_MAP[fn_key]()

四、总结 #

• 如果只是简单的函数对获取需求，字典映射版本更简洁、符合Python的“简单胜于复杂”原则，且完全解决了原代码的冗余问题。
• 如果需要对激活函数进行更多扩展（比如添加属性、自定义方法），类封装版本是更好的选择，结构更清晰，扩展性更强。

内容的提问来源于stack exchange，提问作者Atharva