DFNet with EB3 backbone
| 58 | return features |
| 59 | |
| 60 | class EfficientNetB3(nn.Module): |
| 61 | ''' DFNet with EB3 backbone ''' |
| 62 | default_conf = { |
| 63 | # 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_6"], |
| 64 | 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_5"], |
| 65 | # 'hypercolumn_layers': ["reduction_2", "reduction_4", "reduction_6"], |
| 66 | 'output_dim': 128, |
| 67 | } |
| 68 | mean = [0.485, 0.456, 0.406] |
| 69 | std = [0.229, 0.224, 0.225] |
| 70 | |
| 71 | def __init__(self, feat_dim=12, places365_model_path=''): |
| 72 | super(EfficientNetB3, self).__init__() |
| 73 | # Initialize architecture |
| 74 | self.backbone_net = EfficientNet.from_pretrained('efficientnet-b3') |
| 75 | self.feature_extractor = self.backbone_net.extract_endpoints |
| 76 | |
| 77 | # self.feature_block_index = [1, 3, 6] # same as the 'hypercolumn_layers' |
| 78 | self.feature_block_index = [1, 3, 5] # same as the 'hypercolumn_layers' |
| 79 | # self.feature_block_index = [2, 4, 6] # same as the 'hypercolumn_layers' |
| 80 | |
| 81 | ## adaptation layers, see off branches from fig.3 in S2DNet paper |
| 82 | self.adaptation_layers = AdaptLayers(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim']) |
| 83 | |
| 84 | # pose regression layers |
| 85 | self.avgpool = nn.AdaptiveAvgPool2d(1) |
| 86 | self.fc_pose = nn.Linear(1536, feat_dim) |
| 87 | |
| 88 | def forward(self, x, return_feature=False, isSingleStream=False, upsampleH=120, upsampleW=213): |
| 89 | ''' |
| 90 | inference DFNet. It can regress camera pose as well as extract intermediate layer features. |
| 91 | :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream |
| 92 | :param return_feature: whether to return features as output |
| 93 | :param isSingleStream: whether it's an single stream inference or siamese network inference |
| 94 | :param upsampleH: feature upsample size H |
| 95 | :param upsampleW: feature upsample size W |
| 96 | :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None |
| 97 | :return predict: [2B, 12] or [B, 12] |
| 98 | ''' |
| 99 | # normalize input data |
| 100 | mean, std = x.new_tensor(self.mean), x.new_tensor(self.std) |
| 101 | x = (x - mean[:, None, None]) / std[:, None, None] |
| 102 | |
| 103 | ### encoder ### |
| 104 | feature_maps = [] |
| 105 | list_x = self.feature_extractor(x) |
| 106 | |
| 107 | x = list_x['reduction_6'] # features to save |
| 108 | for i in self.feature_block_index: |
| 109 | fe = list_x['reduction_'+str(i)].clone() |
| 110 | feature_maps.append(fe) |
| 111 | |
| 112 | ### extract and process intermediate features ### |
| 113 | if return_feature: |
| 114 | feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer |
| 115 | |
| 116 | pdb.set_trace() |
| 117 | if isSingleStream: # not siamese network style inference |