| 12345678910111213141516171819202122232425262728293031323334353637383940414243 |
- from model.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)
|
