| 23 | |
| 24 | |
| 25 | class deepFeatureExtractor_EfficientNet(nn.Module): |
| 26 | def __init__(self, architecture="EfficientNet-B5", lv6=False, lv5=False, lv4=False, lv3=False): |
| 27 | super(deepFeatureExtractor_EfficientNet, self).__init__() |
| 28 | assert architecture in ["EfficientNet-B0", "EfficientNet-B1", "EfficientNet-B2", "EfficientNet-B3", |
| 29 | "EfficientNet-B4", "EfficientNet-B5", "EfficientNet-B6", "EfficientNet-B7"] |
| 30 | |
| 31 | if architecture == "EfficientNet-B0": |
| 32 | self.encoder = geffnet.tf_efficientnet_b0_ns(pretrained=True) |
| 33 | self.dimList = [16, 24, 40, 112, 1280] #5th feature is extracted after conv_head or bn2 |
| 34 | #self.dimList = [16, 24, 40, 112, 320] #5th feature is extracted after blocks[6] |
| 35 | elif architecture == "EfficientNet-B1": |
| 36 | self.encoder = geffnet.tf_efficientnet_b1_ns(pretrained=True) |
| 37 | self.dimList = [16, 24, 40, 112, 1280] #5th feature is extracted after conv_head or bn2 |
| 38 | #self.dimList = [16, 24, 40, 112, 320] #5th feature is extracted after blocks[6] |
| 39 | elif architecture == "EfficientNet-B2": |
| 40 | self.encoder = geffnet.tf_efficientnet_b2_ns(pretrained=True) |
| 41 | self.dimList = [16, 24, 48, 120, 1408] #5th feature is extracted after conv_head or bn2 |
| 42 | #self.dimList = [16, 24, 48, 120, 352] #5th feature is extracted after blocks[6] |
| 43 | elif architecture == "EfficientNet-B3": |
| 44 | self.encoder = geffnet.tf_efficientnet_b3_ns(pretrained=True) |
| 45 | self.dimList = [24, 32, 48, 136, 1536] #5th feature is extracted after conv_head or bn2 |
| 46 | #self.dimList = [24, 32, 48, 136, 384] #5th feature is extracted after blocks[6] |
| 47 | elif architecture == "EfficientNet-B4": |
| 48 | self.encoder = geffnet.tf_efficientnet_b4_ns(pretrained=True) |
| 49 | self.dimList = [24, 32, 56, 160, 1792] #5th feature is extracted after conv_head or bn2 |
| 50 | #self.dimList = [24, 32, 56, 160, 448] #5th feature is extracted after blocks[6] |
| 51 | elif architecture == "EfficientNet-B5": |
| 52 | self.encoder = geffnet.tf_efficientnet_b5_ns(pretrained=True) |
| 53 | self.dimList = [24, 40, 64, 176, 2048] #5th feature is extracted after conv_head or bn2 |
| 54 | #self.dimList = [24, 40, 64, 176, 512] #5th feature is extracted after blocks[6] |
| 55 | elif architecture == "EfficientNet-B6": |
| 56 | self.encoder = geffnet.tf_efficientnet_b6_ns(pretrained=True) |
| 57 | self.dimList = [32, 40, 72, 200, 2304] #5th feature is extracted after conv_head or bn2 |
| 58 | #self.dimList = [32, 40, 72, 200, 576] #5th feature is extracted after blocks[6] |
| 59 | elif architecture == "EfficientNet-B7": |
| 60 | self.encoder = geffnet.tf_efficientnet_b7_ns(pretrained=True) |
| 61 | self.dimList = [32, 48, 80, 224, 2560] #5th feature is extracted after conv_head or bn2 |
| 62 | #self.dimList = [32, 48, 80, 224, 640] #5th feature is extracted after blocks[6] |
| 63 | del self.encoder.global_pool |
| 64 | del self.encoder.classifier |
| 65 | #self.block_idx = [3, 4, 5, 7, 9] #5th feature is extracted after blocks[6] |
| 66 | #self.block_idx = [3, 4, 5, 7, 10] #5th feature is extracted after conv_head |
| 67 | self.block_idx = [3, 4, 5, 7, 11] #5th feature is extracted after bn2 |
| 68 | if lv6 is False: |
| 69 | del self.encoder.blocks[6] |
| 70 | del self.encoder.conv_head |
| 71 | del self.encoder.bn2 |
| 72 | del self.encoder.act2 |
| 73 | self.block_idx = self.block_idx[:4] |
| 74 | self.dimList = self.dimList[:4] |
| 75 | if lv5 is False: |
| 76 | del self.encoder.blocks[5] |
| 77 | self.block_idx = self.block_idx[:3] |
| 78 | self.dimList = self.dimList[:3] |
| 79 | if lv4 is False: |
| 80 | del self.encoder.blocks[4] |
| 81 | self.block_idx = self.block_idx[:2] |
| 82 | self.dimList = self.dimList[:2] |
nothing calls this directly
no outgoing calls
no test coverage detected