| 456 | model ref: https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py |
| 457 | ''' |
| 458 | def __init__(self, feat_dim=12, feature_block=6): |
| 459 | super(EfficientNetB3, self).__init__() |
| 460 | self.backbone_net = EfficientNet.from_pretrained('efficientnet-b3') |
| 461 | self.feature_block = feature_block # determine which block's feature to use, max=6 |
| 462 | if self.feature_block == 6: |
| 463 | self.feature_extractor = self.backbone_net.extract_features |
| 464 | else: |
| 465 | self.feature_extractor = self.backbone_net.extract_endpoints |
| 466 | |
| 467 | # self.feature_extractor = self.backbone_net.extract_endpoints # it can restore middle layer |
| 468 | self.avgpool = nn.AdaptiveAvgPool2d(1) |
| 469 | self.fc_pose = nn.Linear(1536, feat_dim) # 1280 for efficientnet-b0, 1536 for efficientnet-b3 |
| 470 | |
| 471 | def _aggregate_feature2(self, x): |
| 472 | ''' |