Added B4 and B5 models

Author: lukemelas

.gitignore

@@ -112,6 +112,8 @@ venv.bak/
 tensorflow/
 example/test*
 *.pth*
+examples/imagenet/data/
+!examples/imagenet/data/README.md
 
 
 

README.md

@@ -1,4 +1,15 @@
 # EfficientNet PyTorch
+
+### Update (June 18, 2019)
+
+The B4 and B5 models are now available. Their usage is identical to the other models: 
+```python
+from efficientnet_pytorch import EfficientNet
+model = EfficientNet.from_pretrained('efficientnet-b4') 
+```
+Upgrade the pip package with `pip install --upgrade pytorch_efficientnet`.
+
+### Overview
 This repository contains an op-for-op PyTorch reimplementation of [EfficientNet](https://arxiv.org/abs/1905.11946), along with pre-trained models and examples. 
 
 The goal of this implementation is to be simple, highly extensible, and easy to integrate into your own projects. This implementation is a work in progress -- new features are currently being implemented. 

efficientnet_pytorch/utils.py

@@ -236,6 +236,8 @@ def get_model_params(model_name, override_params):
     'efficientnet-b1': 'http://storage.googleapis.com/public-models/efficientnet-b1-dbc7070a.pth',
     'efficientnet-b2': 'http://storage.googleapis.com/public-models/efficientnet-b2-27687264.pth',
     'efficientnet-b3': 'http://storage.googleapis.com/public-models/efficientnet-b3-c8376fa2.pth',
+    'efficientnet-b4': 'http://storage.googleapis.com/public-models/efficientnet-b4-e116e8b3.pth',
+    'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet-b5-586e6cc6.pth',
 }
 
 def load_pretrained_weights(model, model_name):

examples/imagenet/main.py

@@ -132,7 +132,11 @@ def main_worker(gpu, ngpus_per_node, args):
         dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                 world_size=args.world_size, rank=args.rank)
     # create model
-    if 'efficientnet' in args.arch:  # NEW
+    if args.arch == 'efficientnet-b5':
+        model = EfficientNet.from_name(args.arch)
+        model.load_state_dict(torch.load('../../tf_to_pytorch/pretrained_pytorch/efficientnet-b5.pth'))
+        print("Using pretrained b5")
+    elif 'efficientnet' in args.arch:  # NEW
         if args.pretrained:
             model = EfficientNet.from_pretrained(args.arch)
             print("=> using pre-trained model '{}'".format(args.arch))

tf_to_pytorch/convert_tf_to_pt/original_tf/efficientnet_model.py

@@ -30,7 +30,8 @@
 from six.moves import xrange  # pylint: disable=redefined-builtin
 import tensorflow as tf
 
-from efficientnet_pytorch import utils
+#from efficientnet_pytorch import utils
+from original_tf import utils
 
 GlobalParams = collections.namedtuple('GlobalParams', [
     'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate', 'data_format',