残差网络中使用门控跳跃连接 代码

门控跳跃连接（Gate Skip Connection）是残差网络中一种常用的连接方式，它能够有效地缓解梯度消失问题，提高网络的训练效果。在门控跳跃连接中，每个残差块都包含两个分支：一个是主分支，另一个是跳跃分支。主分支中的卷积层用于提取特征，跳跃分支中的全连接层用于调整特征通道数，两者通过门控机制相加得到最终的输出。

下面是使用门控跳跃连接的残差块的代码实现：

import torch.nn as nn

class GateSkipBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(GateSkipBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.fc = nn.Linear(in_channels, out_channels, bias=False)
        self.gate = nn.Sigmoid()
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = out + self.gate(self.fc(identity)).unsqueeze(-1).unsqueeze(-1) * identity
        return out

在上述代码中，GateSkipBlock 类继承自 nn.Module 类，实现了一个带有门控跳跃连接的残差块。在 forward 函数中，首先将输入的 x 保存在 identity 变量中，然后在主分支中进行两次卷积操作，并通过门控机制调整跳跃分支的通道数，最后将主分支和跳跃分支相加得到最终的输出。