# add model to __all__ in module model_name = fn.__name__ ifhasattr(mod, '__all__'): mod.__all__.append(model_name) else: mod.__all__ = [model_name]
# add entries to registry dict/sets _model_entrypoints[model_name] = fn _model_to_module[model_name] = module_name _module_to_models[module_name].add(model_name) has_pretrained = False# check if model has a pretrained url to allow filtering on this ifhasattr(mod, 'default_cfgs') and model_name in mod.default_cfgs: # this will catch all models that have entrypoint matching cfg key, but miss any aliasing # entrypoints or non-matching combos has_pretrained = 'url'in mod.default_cfgs[model_name] and'http'in mod.default_cfgs[model_name]['url'] if has_pretrained: _model_has_pretrained.add(model_name) return fn
if kwargs.pop('features_only', False): features = True feature_cfg.setdefault('out_indices', (0, 1, 2, 3, 4)) if'out_indices'in kwargs: feature_cfg['out_indices'] = kwargs.pop('out_indices')
model = model_cls(**kwargs) if model_cfg isNoneelse model_cls(cfg=model_cfg, **kwargs) model.default_cfg = deepcopy(default_cfg)
if pruned: model = adapt_model_from_file(model, variant)
# for classification models, check class attr, then kwargs, then default to 1k, otherwise 0 for feats num_classes_pretrained = 0if features elsegetattr(model, 'num_classes', kwargs.get('num_classes', 1000)) if pretrained: if pretrained_custom_load: load_custom_pretrained(model) else: load_pretrained( model, num_classes=num_classes_pretrained, in_chans=kwargs.get('in_chans', 3), filter_fn=pretrained_filter_fn, strict=pretrained_strict)
if features: feature_cls = FeatureListNet if'feature_cls'in feature_cfg: feature_cls = feature_cfg.pop('feature_cls') ifisinstance(feature_cls, str): feature_cls = feature_cls.lower() if'hook'in feature_cls: feature_cls = FeatureHookNet else: assertFalse, f'Unknown feature class {feature_cls}' model = feature_cls(model, **feature_cfg) model.default_cfg = default_cfg_for_features(default_cfg) # add back default_cfg