| |
| |
| from torch import nn |
| from ..learner import model_meta |
| from ...core import * |
|
|
| pretrainedmodels = try_import('pretrainedmodels') |
| if not pretrainedmodels: |
| raise Exception('Error: `pretrainedmodels` is needed. `pip install pretrainedmodels`') |
|
|
| __all__ = ['inceptionv4', 'inceptionresnetv2', 'nasnetamobile', 'dpn92', 'xception_cadene', 'se_resnet50', |
| 'se_resnet101', 'se_resnext50_32x4d', 'senet154', 'pnasnet5large', 'se_resnext101_32x4d'] |
|
|
| def get_model(model_name:str, pretrained:bool, seq:bool=False, pname:str='imagenet', **kwargs): |
| pretrained = pname if pretrained else None |
| model = getattr(pretrainedmodels, model_name)(pretrained=pretrained, **kwargs) |
| return nn.Sequential(*model.children()) if seq else model |
|
|
| def inceptionv4(pretrained:bool=False): |
| model = get_model('inceptionv4', pretrained) |
| all_layers = list(model.children()) |
| return nn.Sequential(*all_layers[0], *all_layers[1:]) |
| model_meta[inceptionv4] = {'cut': -2, 'split': lambda m: (m[0][11], m[1])} |
|
|
| def nasnetamobile(pretrained:bool=False): |
| model = get_model('nasnetamobile', pretrained, num_classes=1000) |
| model.logits = noop |
| return nn.Sequential(model) |
| model_meta[nasnetamobile] = {'cut': noop, 'split': lambda m: (list(m[0][0].children())[8], m[1])} |
|
|
| def pnasnet5large(pretrained:bool=False): |
| model = get_model('pnasnet5large', pretrained, num_classes=1000) |
| model.logits = noop |
| return nn.Sequential(model) |
| model_meta[pnasnet5large] = {'cut': noop, 'split': lambda m: (list(m[0][0].children())[8], m[1])} |
|
|
| def inceptionresnetv2(pretrained:bool=False): return get_model('inceptionresnetv2', pretrained, seq=True) |
| def dpn92(pretrained:bool=False): return get_model('dpn92', pretrained, pname='imagenet+5k', seq=True) |
| def xception_cadene(pretrained=False): return get_model('xception', pretrained, seq=True) |
| def se_resnet50(pretrained:bool=False): return get_model('se_resnet50', pretrained) |
| def se_resnet101(pretrained:bool=False): return get_model('se_resnet101', pretrained) |
| def se_resnext50_32x4d(pretrained:bool=False): return get_model('se_resnext50_32x4d', pretrained) |
| def se_resnext101_32x4d(pretrained:bool=False): return get_model('se_resnext101_32x4d', pretrained) |
| def senet154(pretrained:bool=False): return get_model('senet154', pretrained) |
|
|
| model_meta[inceptionresnetv2] = {'cut': -2, 'split': lambda m: (m[0][9], m[1])} |
| model_meta[dpn92] = {'cut': -1, 'split': lambda m: (m[0][0][16], m[1])} |
| model_meta[xception_cadene] = {'cut': -1, 'split': lambda m: (m[0][11], m[1])} |
| model_meta[senet154] = {'cut': -3, 'split': lambda m: (m[0][3], m[1])} |
| _se_resnet_meta = {'cut': -2, 'split': lambda m: (m[0][3], m[1])} |
| model_meta[se_resnet50] = _se_resnet_meta |
| model_meta[se_resnet101] = _se_resnet_meta |
| model_meta[se_resnext50_32x4d] = _se_resnet_meta |
| model_meta[se_resnext101_32x4d] = _se_resnet_meta |
|
|
| |
| |
|
|
