def __init__(self, model_config: AttrDict):
super(FeatureExtractorModel, self).__init__()
logging.info("Creating Feature extractor trunk...")
self.model_config = model_config
trunk_name = model_config["TRUNK"]["NAME"]
self.base_model = get_model_trunk(trunk_name)(self.model_config, trunk_name)
self.feature_pool_ops = self._attach_feature_pool_layers()
self._freeze_model()
def _get_trunk(self):
"""
Construct the model trunk given the architecture specified
The trunks could be convnet (AlexNet, ResNe(X)t, RegNet,...etc), transformers etc.
"""
# if we are going to evaluate trunk only we shift to feature extractor backbone
if is_feature_extractor_model(self.model_config):
self.eval_mode = True
return FeatureExtractorModel(self.model_config)
else:
self.eval_mode = False
trunk_name = self.model_config.TRUNK.NAME
return get_model_trunk(trunk_name)(self.model_config, trunk_name)