如何在PyTorch中实现带shared_axes=[1,2]的tf.keras.layers.PReLU？

首先明确两者的核心行为差异：

• TensorFlow中tf.keras.layers.PReLU(shared_axes=[1,2])的作用：
  假设输入为NHWC格式（(batch, H, W, C)），shared_axes=[1,2]指定空间维度H、W共享可学习参数α，即每个通道C对应一个独立的α，同一通道内所有H×W位置共用这个α，参数α的形状为(1,1,1,C)。

• 你提到的nn.PReLU(num_parameters=1)的问题：
  PyTorch的nn.PReLU(num_parameters=1)会让所有通道共享同一个α参数，参数形状为(1)，这和TensorFlow的行为完全不符——TensorFlow是每个通道单独一个α，而这里是全局共用一个。

正确的PyTorch实现 #

要匹配TensorFlow中shared_axes=[1,2]的特性，需要让PyTorch的PReLU为每个通道分配独立的α参数，同时自动在空间维度（H、W）共享该参数（PyTorch会自动完成广播，无需额外配置）。

实现方式是将num_parameters设置为输入的通道数C：

固定通道数的场景 #

import torch.nn as nn

# 假设输入通道数为64
prelu_layer = nn.PReLU(num_parameters=64)

# 输入x为NCHW格式（PyTorch默认）：(batch, C, H, W)
x = torch.randn(32, 64, 28, 28)
output = prelu_layer(x)

动态通道数的场景（自定义模型中） #

如果通道数需要在模型初始化时确定：

import torch.nn as nn

class CustomModel(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        # 根据输入通道数初始化PReLU
        self.prelu = nn.PReLU(num_parameters=in_channels)
    
    def forward(self, x):
        return self.prelu(x)

验证说明 #

PyTorch中nn.PReLU(num_parameters=C)的参数α形状为(C)，在正向传播时会自动广播到NCHW张量的H、W维度，完全等价于TensorFlow中shared_axes=[1,2]的空间共享特性，同时保持每个通道独立的可学习α参数。

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