class AffineLayer(nn.Module): def __init__(self, num_channels, bias=False): super(AffineLayer, self).__init__() weight = torch.FloatTensor(1, num_channels, 1, 1).fill_(1) self.weight = nn.Parameter(weight, requires_grad=True) self.bias = None if bias: bias = torch.FloatTensor(1, num_channels, 1, 1).fill_(0) self.bias = nn.Parameter(bias, requires_grad=True) def forward(self, X): out = X * self.weight.expand_as(X) if self.bias is not None: out = out + self.bias.expand_as(X) return out

这是一个 PyTorch 中的仿射层（Affine Layer）的实现，它可以对输入进行线性变换和平移。其中，num_channels 表示输入的通道数，bias 表示是否使用偏置。在初始化时，权重矩阵的元素都被初始化为 1，偏置矩阵的元素被初始化为 0。在前向传播时，对输入进行线性变换后加上偏置（如果有）。最后返回处理后的输出。