| 131 | |
| 132 | |
| 133 | class MBConv(nn.Module): |
| 134 | def __init__(self, |
| 135 | kernel_size: int, |
| 136 | input_c: int, |
| 137 | out_c: int, |
| 138 | expand_ratio: int, |
| 139 | stride: int, |
| 140 | se_ratio: float, |
| 141 | drop_rate: float, |
| 142 | norm_layer: Callable[..., nn.Module]): |
| 143 | super(MBConv, self).__init__() |
| 144 | |
| 145 | if stride not in [1, 2]: |
| 146 | raise ValueError("illegal stride value.") |
| 147 | |
| 148 | self.has_shortcut = (stride == 1 and input_c == out_c) |
| 149 | |
| 150 | activation_layer = nn.SiLU # alias Swish |
| 151 | expanded_c = input_c * expand_ratio |
| 152 | |
| 153 | # 在EfficientNetV2中，MBConv中不存在expansion=1的情况所以conv_pw肯定存在 |
| 154 | assert expand_ratio != 1 |
| 155 | # Point-wise expansion |
| 156 | self.expand_conv = ConvBNAct(input_c, |
| 157 | expanded_c, |
| 158 | kernel_size=1, |
| 159 | norm_layer=norm_layer, |
| 160 | activation_layer=activation_layer) |
| 161 | |
| 162 | # Depth-wise convolution |
| 163 | self.dwconv = ConvBNAct(expanded_c, |
| 164 | expanded_c, |
| 165 | kernel_size=kernel_size, |
| 166 | stride=stride, |
| 167 | groups=expanded_c, |
| 168 | norm_layer=norm_layer, |
| 169 | activation_layer=activation_layer) |
| 170 | |
| 171 | self.se = SqueezeExcite(input_c, expanded_c, se_ratio) if se_ratio > 0 else nn.Identity() |
| 172 | |
| 173 | # Point-wise linear projection |
| 174 | self.project_conv = ConvBNAct(expanded_c, |
| 175 | out_planes=out_c, |
| 176 | kernel_size=1, |
| 177 | norm_layer=norm_layer, |
| 178 | activation_layer=nn.Identity) # 注意这里没有激活函数，所有传入Identity |
| 179 | |
| 180 | self.out_channels = out_c |
| 181 | |
| 182 | # 只有在使用shortcut连接时才使用dropout层 |
| 183 | self.drop_rate = drop_rate |
| 184 | if self.has_shortcut and drop_rate > 0: |
| 185 | self.dropout = DropPath(drop_rate) |
| 186 | |
| 187 | def forward(self, x: Tensor) -> Tensor: |
| 188 | result = self.expand_conv(x) |
| 189 | result = self.dwconv(result) |
| 190 | result = self.se(result) |
nothing calls this directly
no outgoing calls
no test coverage detected