PyTorch中自定义激活函数的可训练参数如何初始化？

核心结论 #

自定义激活函数的可训练参数应该和单个数据元素的特征维度（即输入的最后一维）一致，而非整个数据集的形状。神经网络处理批量数据时会通过广播机制自动适配批量维度，参数只需对应每个特征的变换逻辑即可。

你的代码合理性分析 #

你当前的代码方向是正确的：

• 参数 units 对应输入数据的特征维度（例如输入形状为 (batch_size, features)，则 units 取值为 features）
• 初始化的 p1、p2、b1、b2 形状均为 (units,)，恰好匹配单个样本的特征数量，PyTorch会自动完成参数与批量输入的广播运算。

关键细节说明 #

• 批量输入适配逻辑：假设输入 inputs 形状为 (batch_size, units)，参数 p1 形状为 (units,)，运算时PyTorch会自动将参数扩展为 (batch_size, units)，与输入执行逐元素运算，无需手动处理批量维度。
• 参数初始化建议：
  • 对于缩放类参数（如你的 p1、p2），用 torch.ones 初始化是合理选择，初始状态下不会改变输入的缩放关系；也可使用小随机值（如 torch.randn(units) * 0.01），避免初始参数过大引发梯度消失或爆炸问题。
  • 偏置项（b1、b2）用 torch.zeros 初始化是行业常规操作，确保初始时不会给输入添加额外偏移。
• 激活函数实现要求：myCustomActivationFunction 需支持广播运算，比如使用逐元素的加减乘除、torch.mul 而非矩阵乘法，这样才能让单特征维度的参数适配批量输入场景。

示例验证 #

假设输入批量大小 batch_size=32，特征维度 units=128，输入形状为 (32, 128)，你的参数 p1 形状为 (128,)，运算时PyTorch会自动将其广播为 (32, 128)，与输入完成逐元素运算，完全满足批量处理需求。

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