You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
95 lines
3.2 KiB
95 lines
3.2 KiB
import sys
|
|
import re
|
|
import fnmatch
|
|
from collections import defaultdict
|
|
|
|
__all__ = ['list_models', 'is_model', 'model_entrypoint', 'list_modules', 'is_model_in_modules']
|
|
|
|
_module_to_models = defaultdict(set) # dict of sets to check membership of model in module
|
|
_model_to_module = {} # mapping of model names to module names
|
|
_model_entrypoints = {} # mapping of model names to entrypoint fns
|
|
_model_has_pretrained = set() # set of model names that have pretrained weight url present
|
|
|
|
|
|
def register_model(fn):
|
|
# lookup containing module
|
|
mod = sys.modules[fn.__module__]
|
|
module_name_split = fn.__module__.split('.')
|
|
module_name = module_name_split[-1] if len(module_name_split) else ''
|
|
|
|
# add model to __all__ in module
|
|
model_name = fn.__name__
|
|
if hasattr(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
|
|
if hasattr(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
|
|
|
|
|
|
def _natural_key(string_):
|
|
return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
|
|
|
|
|
|
def list_models(filter='', module='', pretrained=False):
|
|
""" Return list of available model names, sorted alphabetically
|
|
|
|
Args:
|
|
filter (str) - Wildcard filter string that works with fnmatch
|
|
module (str) - Limit model selection to a specific sub-module (ie 'gen_efficientnet')
|
|
|
|
Example:
|
|
model_list('gluon_resnet*') -- returns all models starting with 'gluon_resnet'
|
|
model_list('*resnext*, 'resnet') -- returns all models with 'resnext' in 'resnet' module
|
|
"""
|
|
if module:
|
|
models = list(_module_to_models[module])
|
|
else:
|
|
models = _model_entrypoints.keys()
|
|
if filter:
|
|
models = fnmatch.filter(models, filter)
|
|
if pretrained:
|
|
models = _model_has_pretrained.intersection(models)
|
|
return list(sorted(models, key=_natural_key))
|
|
|
|
|
|
def is_model(model_name):
|
|
""" Check if a model name exists
|
|
"""
|
|
return model_name in _model_entrypoints
|
|
|
|
|
|
def model_entrypoint(model_name):
|
|
"""Fetch a model entrypoint for specified model name
|
|
"""
|
|
return _model_entrypoints[model_name]
|
|
|
|
|
|
def list_modules():
|
|
""" Return list of module names that contain models / model entrypoints
|
|
"""
|
|
modules = _module_to_models.keys()
|
|
return list(sorted(modules))
|
|
|
|
|
|
def is_model_in_modules(model_name, module_names):
|
|
"""Check if a model exists within a subset of modules
|
|
Args:
|
|
model_name (str) - name of model to check
|
|
module_names (tuple, list, set) - names of modules to search in
|
|
"""
|
|
assert isinstance(module_names, (tuple, list, set))
|
|
return any(model_name in _module_to_models[n] for n in module_names)
|
|
|