1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98
| class EfficientNet(nn.Module): """EfficientNet主网络""" def __init__(self, config, num_classes=1000, include_top=True): super().__init__() self.config = config out_channels = self._round_filters(32) self.stem = nn.Sequential( nn.Conv2d(3, out_channels, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01), SwishActivation() ) blocks = self._make_blocks() self.blocks = nn.Sequential(*blocks) in_channels = blocks[-1].out_channels out_channels = self._round_filters(1280) self.head = nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01), SwishActivation(), nn.AdaptiveAvgPool2d(1) ) if include_top: self.classifier = nn.Sequential( nn.Dropout(config.dropout_rate), nn.Linear(out_channels, num_classes) ) self._initialize_weights() def _round_filters(self, filters): """根据宽度系数缩放滤波器数""" filters *= self.config.width_coef new_filters = max(1, int(filters + self.config.width_coef / 2)) new_filters = int(new_filters / 8) * 8 return new_filters def _make_blocks(self): """构建所有MBConv块""" blocks = [] in_channels = self._round_filters(32) block_configs = [ (1, 16, 1, 1, 3), (6, 24, 2, 2, 3), (6, 40, 2, 2, 5), (6, 80, 3, 2, 3), (6, 112, 3, 1, 5), (6, 192, 4, 2, 5), (6, 320, 1, 1, 3), ] for expand_ratio, out_ch, num_repeat, stride, kernel_size in block_configs: out_channels = self._round_filters(out_ch) strides = [stride] + [1] * (num_repeat - 1) for s in strides: blocks.append(MBConvBlock( in_channels, out_channels, kernel_size, s, expand_ratio )) in_channels = out_channels return blocks def forward(self, x): x = self.stem(x) x = self.blocks(x) x = self.head(x) x = torch.flatten(x, 1) if hasattr(self, 'classifier'): x = self.classifier(x) return x def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out') if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm2d): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) if m.bias is not None: nn.init.zeros_(m.bias)
|