class FactorizedReduce(nn.Module): def __init__(self, C_in, C_out, affine=True): super().__init__() self.relu = nn.ReLU() self.conv1 = nn.Conv2d(C_in, C_out // 2, 1, stride=2, padding=0, bias=False) self.conv2 = nn.Conv2d(C_in, C_out // 2, 1, stride=2, padding=0, bias=False) self.bn = nn.BatchNorm2d(C_out, affine=affine) def forward(self, x): x = self.relu(x) out = torch.cat([self.conv1(x), self.conv2(x[:, :, 1:, 1:])], dim=1) out = self.bn(out) return out
这段代码是一个 PyTorch 的 nn.Module 类,实现了一个名为 FactorizedReduce 的模块。这个模块的作用是将输入的特征图降低尺寸,同时减少通道数,以便在深度神经网络中减少计算量和参数数量。
具体来说,这个模块接收一个 C_in 通道数的特征图作为输入,经过一个 ReLU 激活函数后,分别经过两个 1x1 的卷积层(depthwise convolution),得到两个 C_out/2 通道数的特征图,然后将这两个特征图按通道拼接起来(concatenate),得到一个 C_out 通道数的特征图。最后再经过一个 BatchNorm2d 归一化层,输出降维后的特征图。
值得注意的是,在第二个卷积层中,输入的特征图经过了一个裁剪操作,去掉了左上角的一个像素,这是为了让两个卷积层的输出在空间维度上能够拼接(concatenate)。
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。在前向传播时,对输入进行线性变换后加上偏置(如果有)。最后返回处理后的输出。
class GradientDecoupleLayer(Function): @staticmethod def forward(ctx, x, _lambda): ctx._lambda = _lambda return x @staticmethod def backward(ctx, grad_output): grad_output = grad_output * ctx._lambda return grad_output, None 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 def decouple_layer(x, _lambda): return GradientDecoupleLayer.apply(x, _lambda)
这段代码看起来是 PyTorch 的代码,其中定义了两个类:GradientDecoupleLayer 和 AffineLayer。GradientDecoupleLayer 是一个自定义的 PyTorch 函数,实现了一个梯度解耦的功能;AffineLayer 是一个继承自 nn.Module 的类,实现了一个仿射变换的功能。其中,AffineLayer 的实现比较简单,就是将输入 X 乘上一个可学习的权重参数 self.weight,再加上一个可选的偏置参数 self.bias;而 GradientDecoupleLayer 则是比较复杂的一个函数,它的作用是将输入 x 的梯度乘上一个常数 _lambda,然后输出。这个函数的实现用到了 PyTorch 的静态方法 apply,并且在 backward 函数中返回了梯度 grad_output 乘上 _lambda 的结果。
相关推荐
















