| from efficientnet_pytorch import EfficientNet as _EfficientNet | |
| dependencies = ['torch'] | |
| def _create_model_fn(model_name): | |
| def _model_fn(num_classes=1000, in_channels=3, pretrained='imagenet'): | |
| """Create Efficient Net. | |
| Described in detail here: https://arxiv.org/abs/1905.11946 | |
| Args: | |
| num_classes (int, optional): Number of classes, default is 1000. | |
| in_channels (int, optional): Number of input channels, default | |
| is 3. | |
| pretrained (str, optional): One of [None, 'imagenet', 'advprop'] | |
| If None, no pretrained model is loaded. | |
| If 'imagenet', models trained on imagenet dataset are loaded. | |
| If 'advprop', models trained using adversarial training called | |
| advprop are loaded. It is important to note that the | |
| preprocessing required for the advprop pretrained models is | |
| slightly different from normal ImageNet preprocessing | |
| """ | |
| model_name_ = model_name.replace('_', '-') | |
| if pretrained is not None: | |
| model = _EfficientNet.from_pretrained( | |
| model_name=model_name_, | |
| advprop=(pretrained == 'advprop'), | |
| num_classes=num_classes, | |
| in_channels=in_channels) | |
| else: | |
| model = _EfficientNet.from_name( | |
| model_name=model_name_, | |
| override_params={'num_classes': num_classes}, | |
| ) | |
| model._change_in_channels(in_channels) | |
| return model | |
| return _model_fn | |
| for model_name in ['efficientnet_b' + str(i) for i in range(9)]: | |
| locals()[model_name] = _create_model_fn(model_name) | |