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.
pytorch-image-models/timm/models/factory.py

46 lines
1.6 KiB

from .registry import is_model, is_model_in_modules, model_entrypoint
from .helpers import load_checkpoint
def create_model(
model_name,
pretrained=False,
num_classes=1000,
in_chans=3,
checkpoint_path='',
**kwargs):
"""Create a model
Args:
model_name (str): name of model to instantiate
pretrained (bool): load pretrained ImageNet-1k weights if true
num_classes (int): number of classes for final fully connected layer (default: 1000)
in_chans (int): number of input channels / colors (default: 3)
checkpoint_path (str): path of checkpoint to load after model is initialized
Keyword Args:
drop_rate (float): dropout rate for training (default: 0.0)
global_pool (str): global pool type (default: 'avg')
**: other kwargs are model specific
"""
margs = dict(pretrained=pretrained, num_classes=num_classes, in_chans=in_chans)
# Only EfficientNet and MobileNetV3 models have support for batchnorm params or drop_connect_rate passed as args
is_efficientnet = is_model_in_modules(model_name, ['efficientnet', 'mobilenetv3'])
if not is_efficientnet:
kwargs.pop('bn_tf', None)
kwargs.pop('bn_momentum', None)
kwargs.pop('bn_eps', None)
kwargs.pop('drop_connect_rate', None)
if is_model(model_name):
create_fn = model_entrypoint(model_name)
model = create_fn(**margs, **kwargs)
else:
raise RuntimeError('Unknown model (%s)' % model_name)
if checkpoint_path:
load_checkpoint(model, checkpoint_path)
return model