Merge pull request #57 from lukemelas/add_b7

Add B6, B7 models and AutoAugment models

Author: lukemelas

README.md

@@ -1,9 +1,17 @@
 # EfficientNet PyTorch
 
-### Update (June 29, 2019)
+### Update (July 31, 2019)
 
 _Upgrade the pip package with_ `pip install --upgrade efficientnet-pytorch`
 
+The B6 and B7 models are now available. Additionally, _all_ pretrained models have been updated to use AutoAugment preprocessing, which translates to better performance across the board. Usage is the same as before: 
+```python
+from efficientnet_pytorch import EfficientNet
+model = EfficientNet.from_pretrained('efficientnet-b7') 
+```
+
+### Update (June 29, 2019)
+
 This update adds easy model exporting ([#20](https://github.com/lukemelas/EfficientNet-PyTorch/issues/20)) and feature extraction ([#38](https://github.com/lukemelas/EfficientNet-PyTorch/issues/38)). 
 
  * [Example: Export to ONNX](#example-export)
@@ -124,8 +132,8 @@ Details about the models are below:
 | `efficientnet-b3` |    12M   |    81.1    |      ✓      |
 | `efficientnet-b4` |    19M   |    82.6    |      ✓      |
 | `efficientnet-b5` |    30M   |    83.3    |      ✓      |
-| `efficientnet-b6` |    43M   |    84.0    |      -      |
-| `efficientnet-b7` |    66M   |    84.4    |      -      |
+| `efficientnet-b6` |    43M   |    84.0    |      ✓      |
+| `efficientnet-b7` |    66M   |    84.4    |      ✓      |
 
 
 #### Example: Classification

efficientnet_pytorch/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 from .model import EfficientNet
 from .utils import (
     GlobalParams,

efficientnet_pytorch/utils.py

@@ -276,12 +276,14 @@ def get_model_params(model_name, override_params):
 
 
 url_map = {
-    'efficientnet-b0': 'http://storage.googleapis.com/public-models/efficientnet-b0-08094119.pth',
-    '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',
+    'efficientnet-b0': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b0-355c32eb.pth',
+    'efficientnet-b1': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b1-f1951068.pth',
+    'efficientnet-b2': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b2-8bb594d6.pth',
+    'efficientnet-b3': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b3-5fb5a3c3.pth',
+    'efficientnet-b4': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b4-6ed6700e.pth',
+    'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b5-b6417697.pth',
+    'efficientnet-b6': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b6-c76e70fd.pth',
+    'efficientnet-b7': 'http://storage.googleapis.com/public-models/efficientnet/efficientnet-b7-dcc49843.pth',
 }
 
 def load_pretrained_weights(model, model_name, load_fc=True):

setup.py

@@ -18,7 +18,7 @@
 EMAIL = '[email protected]'
 AUTHOR = 'Luke'
 REQUIRES_PYTHON = '>=3.5.0'
-VERSION = '0.3.0'
+VERSION = '0.4.0'
 
 # What packages are required for this module to be executed?
 REQUIRED = [

tf_to_pytorch/README.md

@@ -7,7 +7,7 @@ I should also emphasize that you do *not* need to run any of this code to load p
 That being said, the main script here is `convert_to_tf/load_tf_weights.py`. In order to use it, you should first download the pretrained TensorFlow weights:
  ```bash
 cd pretrained_tensorflow
-./download.sh
+./download.sh efficientnet-b0
 cd ..
 ```
 Then
@@ -19,3 +19,7 @@ python load_tf_weights.py \
     --tf_checkpoint ../pretrained_tensorflow/efficientnet-b0/ \
     --output_file ../pretrained_pytorch/efficientnet-b0.pth
 ``` 
+
+<!-- Here is a helpful utility:
+mv efficientnet-b0.pth efficientnet-b0-$(sha256sum efficientnet-b0.pth | head -c 8).pth
+-->

tf_to_pytorch/pretrained_tensorflow/download.sh

@@ -9,6 +9,6 @@
 #   ./download.sh efficientnet-b3
 
 MODEL=$1
-wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz
+wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/${MODEL}.tar.gz
 tar xvf ${MODEL}.tar.gz
 rm ${MODEL}.tar.gz