Merge pull request #214 from chrisyeh96/constant_model_names

Move valid model names to constant

Author: lukemelas

efficientnet_pytorch/__init__.py

@@ -1,5 +1,5 @@
 __version__ = "0.7.0"
-from .model import EfficientNet
+from .model import EfficientNet, VALID_MODELS
 from .utils import (
     GlobalParams,
     BlockArgs,

efficientnet_pytorch/model.py

@@ -22,6 +22,17 @@
     calculate_output_image_size
 )
 
+
+VALID_MODELS = (
+    'efficientnet-b0', 'efficientnet-b1', 'efficientnet-b2', 'efficientnet-b3',
+    'efficientnet-b4', 'efficientnet-b5', 'efficientnet-b6', 'efficientnet-b7',
+    'efficientnet-b8',
+
+    # Support the construction of 'efficientnet-l2' without pretrained weights
+    'efficientnet-l2'
+)
+
+
 class MBConvBlock(nn.Module):
     """Mobile Inverted Residual Bottleneck Block.
 
@@ -388,14 +399,9 @@ def _check_model_name_is_valid(cls, model_name):
         Returns:
             bool: Is a valid name or not.
         """
-        valid_models = ['efficientnet-b'+str(i) for i in range(9)]
-
-        # Support the construction of 'efficientnet-l2' without pretrained weights
-        valid_models += ['efficientnet-l2']
+        if model_name not in VALID_MODELS:
+            raise ValueError('model_name should be one of: ' + ', '.join(VALID_MODELS))
 
-        if model_name not in valid_models:
-            raise ValueError('model_name should be one of: ' + ', '.join(valid_models))
-    
     def _change_in_channels(self, in_channels):
         """Adjust model's first convolution layer to in_channels, if in_channels not equals 3.