Class EfficientNetB0

script/feature/efficientnet.py:181–271

DFNet with EB0 backbone, feature levels can be customized

        return features

class EfficientNetB0(nn.Module):
    ''' DFNet with EB0 backbone, feature levels can be customized '''
    default_conf = {
        # 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_6"],
        'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_5"],
        # 'hypercolumn_layers': ["reduction_2", "reduction_4", "reduction_6"],
        # 'hypercolumn_layers': ["reduction_1"],
        'output_dim': 128,
    }
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]

    def __init__(self, feat_dim=12, places365_model_path=''):
        super(EfficientNetB0, self).__init__()
        # Initialize architecture
        self.backbone_net = EfficientNet.from_pretrained('efficientnet-b0')
        self.feature_extractor = self.backbone_net.extract_endpoints

        # self.feature_block_index = [1, 3, 6] # same as the 'hypercolumn_layers'
        self.feature_block_index = [1, 3, 5] # same as the 'hypercolumn_layers'
        # self.feature_block_index = [2, 4, 6] # same as the 'hypercolumn_layers'
        # self.feature_block_index = [1]

        ## adaptation layers, see off branches from fig.3 in S2DNet paper
        self.adaptation_layers = AdaptLayers2(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim'])

        # pose regression layers
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc_pose = nn.Linear(1280, feat_dim)

    def forward(self, x, return_feature=False, isSingleStream=False, return_pose=False, upsampleH=120, upsampleW=213):
        '''
        inference DFNet. It can regress camera pose as well as extract intermediate layer features.
            :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream
            :param return_feature: whether to return features as output
            :param isSingleStream: whether it's an single stream inference or siamese network inference
            :param return_pose: TODO: if only return_pose, we don't need to compute return_feature part
            :param upsampleH: feature upsample size H
            :param upsampleW: feature upsample size W
            :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None
            :return predict: [2B, 12] or [B, 12]
        '''
        # normalize input data
        mean, std = x.new_tensor(self.mean), x.new_tensor(self.std)
        x = (x - mean[:, None, None]) / std[:, None, None]

        ### encoder ###
        feature_maps = []
        list_x = self.feature_extractor(x)

        x = list_x['reduction_6'] # features to save
        for i in self.feature_block_index:
            fe = list_x['reduction_'+str(i)].clone()
            feature_maps.append(fe)

        ### extract and process intermediate features ###
        if return_feature:
            feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer