Merge pull request #44 from lukemelas/add-export-and-improve-code

Add export and improve code

Author: lukemelas

.gitignore

@@ -114,6 +114,10 @@ example/test*
 *.pth*
 examples/imagenet/data/
 !examples/imagenet/data/README.md
+tmp
+tf_to_pytorch/pretrained_tensorflow
+!tf_to_pytorch/pretrained_tensorflow/download.sh
+examples/imagenet/run.sh
 
 
 

README.md

@@ -1,13 +1,28 @@
 # EfficientNet PyTorch
 
+### Update (June 29, 2019)
+
+_Upgrade the pip package with_ `pip install --upgrade efficientnet-pytorch`
+
+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)
+ * [Example: Extract features](#example-feature-extraction)
+ * Also: fixed a CUDA/CPU bug ([#32](https://github.com/lukemelas/EfficientNet-PyTorch/issues/32))
+
+It is also now incredibly simple to load a pretrained model with a new number of classes for transfer learning:
+```python
+model = EfficientNet.from_pretrained('efficientnet-b1', num_classes=23)
+``` 
+
+
 ### Update (June 23, 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 efficientnet-pytorch`.
 
 ### 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. 
@@ -32,6 +47,7 @@ _Upcoming features_: In the next few days, you will be able to:
     * [Load pretrained models](#loading-pretrained-models)
     * [Example: Classify](#example-classification)
     * [Example: Extract features](#example-feature-extraction)
+    * [Example: Export to ONNX](#example-export)
 6. [Contributing](#contributing) 
 
 ### About EfficientNet
@@ -160,9 +176,25 @@ model = EfficientNet.from_pretrained('efficientnet-b0')
 print(img.shape) # torch.Size([1, 3, 224, 224])
 
 features = model.extract_features(img)
-print(features.shape) # torch.Size([1, 320, 7, 7])
+print(features.shape) # torch.Size([1, 1280, 7, 7])
 ```
 
+#### Example: Export to ONNX  
+
+Exporting to ONNX for deploying to production is now simple: 
+```python
+import torch 
+from efficientnet_pytorch import EfficientNet
+
+model = EfficientNet.from_pretrained('efficientnet-b1')
+dummy_input = torch.randn(10, 3, 240, 240)
+
+torch.onnx.export(model, dummy_input, "test-b1.onnx", verbose=True)
+``` 
+
+[Here](https://colab.research.google.com/drive/1rOAEXeXHaA8uo3aG2YcFDHItlRJMV0VP) is a Colab example. 
+
+
 #### ImageNet
 
 See `examples/imagenet` for details about evaluating on ImageNet.

efficientnet_pytorch/__init__.py

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

efficientnet_pytorch/model.py

@@ -7,7 +7,7 @@
     round_filters,
     round_repeats,
     drop_connect,
-    Conv2dSamePadding,
+    get_same_padding_conv2d,
     get_model_params,
     efficientnet_params,
     load_pretrained_weights,
@@ -33,30 +33,33 @@ def __init__(self, block_args, global_params):
         self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)
         self.id_skip = block_args.id_skip  # skip connection and drop connect
 
+        # Get static or dynamic convolution depending on image size
+        Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
+
         # Expansion phase
         inp = self._block_args.input_filters  # number of input channels
         oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels
         if self._block_args.expand_ratio != 1:
-            self._expand_conv = Conv2dSamePadding(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
+            self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
             self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
 
         # Depthwise convolution phase
         k = self._block_args.kernel_size
         s = self._block_args.stride
-        self._depthwise_conv = Conv2dSamePadding(
+        self._depthwise_conv = Conv2d(
             in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise
             kernel_size=k, stride=s, bias=False)
         self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
 
         # Squeeze and Excitation layer, if desired
         if self.has_se:
             num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
-            self._se_reduce = Conv2dSamePadding(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
-            self._se_expand = Conv2dSamePadding(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)
+            self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
+            self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)
 
         # Output phase
         final_oup = self._block_args.output_filters
-        self._project_conv = Conv2dSamePadding(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
+        self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
         self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)
 
     def forward(self, inputs, drop_connect_rate=None):
@@ -109,14 +112,17 @@ def __init__(self, blocks_args=None, global_params=None):
         self._global_params = global_params
         self._blocks_args = blocks_args
 
+        # Get static or dynamic convolution depending on image size
+        Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
+
         # Batch norm parameters
         bn_mom = 1 - self._global_params.batch_norm_momentum
         bn_eps = self._global_params.batch_norm_epsilon
 
         # Stem
         in_channels = 3  # rgb
         out_channels = round_filters(32, self._global_params)  # number of output channels
-        self._conv_stem = Conv2dSamePadding(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
+        self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
         self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
 
         # Build blocks
@@ -140,7 +146,7 @@ def __init__(self, blocks_args=None, global_params=None):
         # Head
         in_channels = block_args.output_filters  # output of final block
         out_channels = round_filters(1280, self._global_params)
-        self._conv_head = Conv2dSamePadding(in_channels, out_channels, kernel_size=1, bias=False)
+        self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
         self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
 
         # Final linear layer
@@ -158,7 +164,10 @@ def extract_features(self, inputs):
             drop_connect_rate = self._global_params.drop_connect_rate
             if drop_connect_rate:
                 drop_connect_rate *= float(idx) / len(self._blocks)
-            x = block(x, drop_connect_rate) 
+            x = block(x, drop_connect_rate=drop_connect_rate)
+
+        # Head
+        x = relu_fn(self._bn1(self._conv_head(x)))
 
         return x
 
@@ -168,8 +177,7 @@ def forward(self, inputs):
         # Convolution layers
         x = self.extract_features(inputs)
 
-        # Head
-        x = relu_fn(self._bn1(self._conv_head(x)))
+        # Pooling and final linear layer
         x = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1)
         if self._dropout:
             x = F.dropout(x, p=self._dropout, training=self.training)
@@ -183,9 +191,9 @@ def from_name(cls, model_name, override_params=None):
         return EfficientNet(blocks_args, global_params)
 
     @classmethod
-    def from_pretrained(cls, model_name):
-        model = EfficientNet.from_name(model_name)
-        load_pretrained_weights(model, model_name)
+    def from_pretrained(cls, model_name, num_classes=1000):
+        model = EfficientNet.from_name(model_name, override_params={'num_classes': num_classes})
+        load_pretrained_weights(model, model_name, load_fc=(num_classes == 1000))
         return model
 
     @classmethod

efficientnet_pytorch/utils.py

@@ -6,6 +6,7 @@
 import re
 import math
 import collections
+from functools import partial
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -21,7 +22,7 @@
 GlobalParams = collections.namedtuple('GlobalParams', [
     'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate',
     'num_classes', 'width_coefficient', 'depth_coefficient',
-    'depth_divisor', 'min_depth', 'drop_connect_rate',])
+    'depth_divisor', 'min_depth', 'drop_connect_rate', 'image_size'])
 
 
 # Parameters for an individual model block
@@ -75,8 +76,16 @@ def drop_connect(inputs, p, training):
     return output
 
 
-class Conv2dSamePadding(nn.Conv2d):
-    """ 2D Convolutions like TensorFlow """
+def get_same_padding_conv2d(image_size=None):
+    """ Chooses static padding if you have specified an image size, and dynamic padding otherwise.
+        Static padding is necessary for ONNX exporting of models. """
+    if image_size is None:
+        return Conv2dDynamicSamePadding
+    else:
+        return partial(Conv2dStaticSamePadding, image_size=image_size)
+
+class Conv2dDynamicSamePadding(nn.Conv2d):
+    """ 2D Convolutions like TensorFlow, for a dynamic image size """
     def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True):
         super().__init__(in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
         self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]]*2
@@ -93,6 +102,39 @@ def forward(self, x):
         return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
 
 
+class Conv2dStaticSamePadding(nn.Conv2d):
+    """ 2D Convolutions like TensorFlow, for a fixed image size"""
+    def __init__(self, in_channels, out_channels, kernel_size, image_size=None, **kwargs):
+        super().__init__(in_channels, out_channels, kernel_size, **kwargs)
+        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2
+
+        # Calculate padding based on image size and save it
+        assert image_size is not None
+        ih, iw = image_size if type(image_size) == list else [image_size, image_size]
+        kh, kw = self.weight.size()[-2:]
+        sh, sw = self.stride
+        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)
+        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)
+        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)
+        if pad_h > 0 or pad_w > 0:
+            self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2))
+        else:
+            self.static_padding = Identity()
+
+    def forward(self, x):
+        x = self.static_padding(x)
+        x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+        return x
+
+
+class Identity(nn.Module):
+    def __init__(self,):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
 ########################################################################
 ############## HELPERS FUNCTIONS FOR LOADING MODEL PARAMS ##############
 ########################################################################
@@ -189,8 +231,8 @@ def encode(blocks_args):
         return block_strings
 
 
-def efficientnet(width_coefficient=None, depth_coefficient=None,
-                 dropout_rate=0.2, drop_connect_rate=0.2):
+def efficientnet(width_coefficient=None, depth_coefficient=None, dropout_rate=0.2,
+                 drop_connect_rate=0.2, image_size=None, num_classes=1000):
     """ Creates a efficientnet model. """
 
     blocks_args = [
@@ -207,11 +249,12 @@ def efficientnet(width_coefficient=None, depth_coefficient=None,
         dropout_rate=dropout_rate,
         drop_connect_rate=drop_connect_rate,
         # data_format='channels_last',  # removed, this is always true in PyTorch
-        num_classes=1000,
+        num_classes=num_classes,
         width_coefficient=width_coefficient,
         depth_coefficient=depth_coefficient,
         depth_divisor=8,
-        min_depth=None
+        min_depth=None,
+        image_size=image_size,
     )
 
     return blocks_args, global_params
@@ -220,9 +263,10 @@ def efficientnet(width_coefficient=None, depth_coefficient=None,
 def get_model_params(model_name, override_params):
     """ Get the block args and global params for a given model """
     if model_name.startswith('efficientnet'):
-        w, d, _, p = efficientnet_params(model_name)
+        w, d, s, p = efficientnet_params(model_name)
         # note: all models have drop connect rate = 0.2
-        blocks_args, global_params = efficientnet(width_coefficient=w, depth_coefficient=d, dropout_rate=p)
+        blocks_args, global_params = efficientnet(
+            width_coefficient=w, depth_coefficient=d, dropout_rate=p, image_size=s)
     else:
         raise NotImplementedError('model name is not pre-defined: %s' % model_name)
     if override_params:
@@ -240,8 +284,14 @@ def get_model_params(model_name, override_params):
     'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet-b5-586e6cc6.pth',
 }
 
-def load_pretrained_weights(model, model_name):
+def load_pretrained_weights(model, model_name, load_fc=True):
     """ Loads pretrained weights, and downloads if loading for the first time. """
     state_dict = model_zoo.load_url(url_map[model_name])
-    model.load_state_dict(state_dict)
+    if load_fc:
+        model.load_state_dict(state_dict)
+    else:
+        state_dict.pop('_fc.weight')
+        state_dict.pop('_fc.bias')
+        res = model.load_state_dict(state_dict, strict=False)
+        assert str(res.missing_keys) == str(['_fc.weight', '_fc.bias']), 'issue loading pretrained weights'
     print('Loaded pretrained weights for {}'.format(model_name))

setup.py

@@ -18,7 +18,7 @@
 EMAIL = '[email protected]'
 AUTHOR = 'Luke'
 REQUIRES_PYTHON = '>=3.5.0'
-VERSION = '0.2.0'
+VERSION = '0.3.0'
 
 # What packages are required for this module to be executed?
 REQUIRED = [
@@ -109,16 +109,13 @@ def run(self):
     extras_require=EXTRAS,
     include_package_data=True,
     license='Apache',
-    # classifiers=[
-    #     # Trove classifiers
-    #     # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
-    #     'License :: OSI Approved :: MIT License',
-    #     'Programming Language :: Python',
-    #     'Programming Language :: Python :: 3',
-    #     'Programming Language :: Python :: 3.6',
-    #     'Programming Language :: Python :: Implementation :: CPython',
-    #     'Programming Language :: Python :: Implementation :: PyPy'
-    # ],
+    classifiers=[
+        # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.6',
+    ],
     # $ setup.py publish support.
     cmdclass={
         'upload': UploadCommand,