MeliusNet22 implementation (#142)

Koen Helwegen · lgeiger · AdamHillier · web-flow · commit 8cb6eded8544 · 2020-04-24T18:46:08.000+01:00
* MeliusNet22 implementation

* Apply suggestions from code review

Co-Authored-By: Lukas Geiger &lt;lgeiger@users.noreply.github.com&gt;

* review remarks

* Update larq_zoo/literature/meliusnet.py

Co-Authored-By: Lukas Geiger &lt;lgeiger@users.noreply.github.com&gt;

* wip

* correct hashes, docstr etc

* isort

* Apply suggestions from code review

Co-Authored-By: Adam Hillier &lt;7688302+AdamHillier@users.noreply.github.com&gt;
Co-Authored-By: Lukas Geiger &lt;lgeiger@users.noreply.github.com&gt;

* import tuple

* linting

* explicit naming

* add path attributes to base

* better TFOpLayer, update hashes

* inline if/else

Co-authored-by: Lukas Geiger &lt;lgeiger@users.noreply.github.com&gt;
Co-authored-by: Adam Hillier &lt;7688302+AdamHillier@users.noreply.github.com&gt;
diff --git a/larq_zoo/core/utils.py b/larq_zoo/core/utils.py
@@ -25,11 +25,7 @@ def slash_join(*args):
 
 
 def download_pretrained_model(
-    model: str,
-    version: str,
-    file: str,
-    file_hash: str,
-    cache_dir: Optional[str] = None,
+    model: str, version: str, file: str, file_hash: str, cache_dir: Optional[str] = None
 ) -> str:
     root_url = "https://github.com/larq/zoo/releases/download/"
 
@@ -144,10 +140,12 @@ def global_pool(
         pool_size = (
             input_shape[1:3] if data_format == "channels_last" else input_shape[2:4]
         )
-        x = keras.layers.AveragePooling2D(pool_size=pool_size, data_format=data_format)(
-            x
-        )
-        x = keras.layers.Flatten()(x)
+        x = keras.layers.AveragePooling2D(
+            pool_size=pool_size,
+            data_format=data_format,
+            name=f"{name}_pool" if name else None,
+        )(x)
+        x = keras.layers.Flatten(name=f"{name}_flatten" if name else None)(x)
     except ValueError:
         x = keras.layers.GlobalAveragePooling2D(data_format=data_format, name=name)(x)
 
@@ -170,3 +168,20 @@ def decode_predictions(preds, top=5, **kwargs):
     ValueError: In case of invalid shape of the `pred` array (must be 2D).
     """
     return keras_decode_predictions(preds, top=top, **kwargs)
+
+
+def TFOpLayer(tf_op: tf.Operation, *args, **kwargs) -> tf.keras.layers.Layer:
+    """Wrap a tensorflow op using a Lambda layer. This facilitates naming the op as a
+    proper keras layer.
+
+    Example: `TFOpLayer(tf.split, groups, axis=-1, name="split")(x)`.
+
+    # Arguments
+    tf_op: tensorflow that needs to be wrapped.
+
+    # Returns
+    A keras layer wrapping `tf_op`.
+
+    """
+    name = kwargs.pop("name", None)
+    return tf.keras.layers.Lambda(lambda x_: tf_op(x_, *args, **kwargs), name=name)
diff --git a/larq_zoo/literature/__init__.py b/larq_zoo/literature/__init__.py
@@ -7,6 +7,7 @@
     BinaryDenseNet45,
 )
 from larq_zoo.literature.dorefanet import DoReFaNet
+from larq_zoo.literature.meliusnet import MeliusNet22
 from larq_zoo.literature.real_to_bin_nets import RealToBinaryNet
 from larq_zoo.literature.resnet_e import BinaryResNetE18
 from larq_zoo.literature.xnornet import XNORNet
@@ -20,6 +21,7 @@
     "BinaryDenseNet37Dilated",
     "BinaryDenseNet45",
     "DoReFaNet",
+    "MeliusNet22",
     "RealToBinaryNet",
     "XNORNet",
 ]
diff --git a/larq_zoo/literature/meliusnet.py b/larq_zoo/literature/meliusnet.py
@@ -0,0 +1,278 @@
+from typing import Optional, Sequence, Tuple, Union
+
+import larq as lq
+import tensorflow as tf
+from zookeeper import Field, factory
+
+from larq_zoo.core import utils
+from larq_zoo.core.model_factory import ModelFactory
+
+################
+# Base factory #
+################
+
+
+class MeliusNetFactory(ModelFactory):
+    # Overall architecture configuration. These are not `Fields`, as they should
+    # not be configurable, but set in the various concrete subclasses.
+    num_blocks: Sequence[int]
+    transition_features: Sequence[int]
+    name: str = None
+    imagenet_weights_path: str
+    imagenet_no_top_weights_path: str
+
+    # Some default layer arguments.
+    batch_norm_momentum: float = Field(0.9)
+    kernel_initializer: Optional[Union[str, tf.keras.initializers.Initializer]] = Field(
+        "glorot_normal"
+    )
+    input_quantizer = Field(lambda: lq.quantizers.SteSign(1.3))
+    kernel_quantizer = Field(lambda: lq.quantizers.SteSign(1.3))
+    kernel_constraint = Field(lambda: lq.constraints.WeightClip(1.3))
+
+    def pool(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        return tf.keras.layers.MaxPool2D(2, strides=2, padding="same", name=name)(x)
+
+    def norm(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        return tf.keras.layers.BatchNormalization(
+            momentum=self.batch_norm_momentum, epsilon=1e-5, name=name
+        )(x)
+
+    def act(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        return tf.keras.layers.Activation("relu", name=name)(x)
+
+    def quant_conv(
+        self,
+        x: tf.Tensor,
+        filters: int,
+        kernel: Union[int, Tuple[int, int]],
+        strides: Union[int, Tuple[int, int]] = 1,
+        name: str = None,
+    ) -> tf.Tensor:
+        return lq.layers.QuantConv2D(
+            filters,
+            kernel,
+            strides=strides,
+            padding="same",
+            use_bias=False,
+            input_quantizer=self.input_quantizer,
+            kernel_quantizer=self.kernel_quantizer,
+            kernel_constraint=self.kernel_constraint,
+            kernel_initializer=self.kernel_initializer,
+            name=name,
+        )(x)
+
+    def group_conv(
+        self,
+        x: tf.Tensor,
+        filters: int,
+        kernel: Union[int, Tuple[int, int]],
+        groups: int,
+        name: str = None,
+    ) -> tf.Tensor:
+        assert filters % groups == 0
+        assert x.shape.as_list()[-1] % groups == 0
+
+        x_split = utils.TFOpLayer(tf.split, groups, axis=-1, name=f"{name}_split")(x)
+
+        y_split = [
+            tf.keras.layers.Conv2D(
+                filters // groups,
+                kernel,
+                padding="same",
+                use_bias=False,
+                kernel_initializer=self.kernel_initializer,
+                name=f"{name}_conv{i}",
+            )(split)
+            for i, split in enumerate(x_split)
+        ]
+
+        return utils.TFOpLayer(tf.concat, axis=-1, name=f"{name}_concat")(y_split)
+
+    def group_stem(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        x = tf.keras.layers.Conv2D(
+            32,
+            3,
+            strides=2,
+            padding="same",
+            use_bias=False,
+            kernel_initializer=self.kernel_initializer,
+            name=f"{name}_s0_conv",
+        )(x)
+        x = self.norm(x, name=f"{name}_s0_bn")
+        x = self.act(x, name=f"{name}_s0_relu")
+
+        x = self.group_conv(x, 32, 3, 4, name=f"{name}_s1_groupconv")
+        x = self.norm(x, name=f"{name}_s1_bn")
+        x = self.act(x, name=f"{name}_s1_relu")
+
+        x = self.group_conv(x, 64, 3, 8, name=f"{name}_s2_groupconv")
+        x = self.norm(x, name=f"{name}_s2_bn")
+        x = self.act(x, name=f"{name}_s2_relu")
+
+        return self.pool(x, name=f"{name}_pool")
+
+    def dense_block(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        w = x
+        w = self.norm(w, name=f"{name}_bn")
+        w = self.quant_conv(w, 64, 3, name=f"{name}_binconv")
+        return utils.TFOpLayer(tf.concat, axis=-1, name=f"{name}_concat")([x, w])
+
+    def improvement_block(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        w = x
+        w = self.norm(w, name=f"{name}_bn")
+        w = self.quant_conv(w, 64, 3, name=f"{name}_binconv")
+        f_in = int(x.shape[-1])
+        return tf.keras.layers.Lambda(
+            lambda x_: x_[0] + tf.pad(x_[1], [[0, 0], [0, 0], [0, 0], [f_in - 64, 0]]),
+            name=f"{name}_merge",
+        )([x, w])
+
+    def transition_block(
+        self, x: tf.Tensor, filters: int, name: str = None
+    ) -> tf.Tensor:
+        x = self.norm(x, name=f"{name}_bn")
+        x = self.pool(x, name=f"{name}_maxpool")
+        x = self.act(x, name=f"{name}_relu")
+        return tf.keras.layers.Conv2D(
+            filters,
+            1,
+            use_bias=False,
+            kernel_initializer=self.kernel_initializer,
+            name=f"{name}_pw",
+        )(x)
+
+    def block(self, x: tf.Tensor, name: str = None) -> tf.Tensor:
+        x = self.dense_block(x, name=f"{name}_dense")
+        return self.improvement_block(x, name=f"{name}_improve")
+
+    def build(self) -> tf.keras.models.Model:
+        x = self.image_input
+        x = self.group_stem(x, name="stem")
+        for i, (n, f) in enumerate(zip(self.num_blocks, self.transition_features)):
+            for j in range(n):
+                x = self.block(x, f"section_{i}_block_{j}")
+            if f:
+                x = self.transition_block(x, f, f"section_{i}_transition")
+
+        x = self.norm(x, "head_bn")
+        x = self.act(x, "head_relu")
+
+        if self.include_top:
+            x = utils.global_pool(x, name="head_globalpool")
+            x = tf.keras.layers.Dense(
+                self.num_classes,
+                kernel_initializer=self.kernel_initializer,
+                name="head_dense",
+            )(x)
+            x = tf.keras.layers.Activation(
+                "softmax", dtype="float32", name="head_softmax"
+            )(x)
+
+        model = tf.keras.models.Model(
+            inputs=self.image_input, outputs=x, name=self.name
+        )
+
+        if self.weights == "imagenet":
+            model.load_weights(
+                self.imagenet_weights_path
+                if self.include_top
+                else self.imagenet_no_top_weights_path
+            )
+        elif self.weights is not None:
+            model.load_weights(self.weights)
+
+        return model
+
+
+######################
+# Concrete factories #
+######################
+
+
+@factory
+class MeliusNet22Factory(MeliusNetFactory):
+    num_blocks = (4, 5, 4, 4)
+    transition_features = (160, 224, 256, None)
+    name = "meliusnet22"
+
+    @property
+    def imagenet_weights_path(self):
+        return utils.download_pretrained_model(
+            model="meliusnet22",
+            version="v0.1.0",
+            file="meliusnet22_weights.h5",
+            file_hash="c1ba85e8389ae326009665ec13331e49fc3df4d0f925fa8553e224f7362c18ed",
+        )
+
+    @property
+    def imagenet_no_top_weights_path(self):
+        return utils.download_pretrained_model(
+            model="meliusnet22",
+            version="v0.1.0",
+            file="meliusnet22_weights_notop.h5",
+            file_hash="b64c8296a3d07ce2799846caf0ad6d390f6cd9bbf21ea3390fafbab87bb79aa5",
+        )
+
+
+#########################
+# Functional interfaces #
+#########################
+
+
+def MeliusNet22(
+    *,  # Keyword arguments only
+    input_shape: Optional[Sequence[Optional[int]]] = None,
+    input_tensor: Optional[tf.Tensor] = None,
+    weights: Optional[str] = "imagenet",
+    include_top: bool = True,
+    num_classes: int = 1000,
+) -> tf.keras.models.Model:
+    """Instantiates the MeliusNet22 architecture.
+
+    Optionally loads weights pre-trained on ImageNet.
+
+    ```netron
+    meliusnet22-v0.1.0/meliusnet22.json
+    ```
+    ```summary
+    literature.MeliusNet22
+    ```
+    ```plot-altair
+    /plots/meliusnet22.vg.json
+    ```
+
+    # ImageNet Metrics
+
+    | Top-1 Accuracy | Top-5 Accuracy | Parameters | Memory   |
+    | -------------- | -------------- | ---------- | -------- |
+    | 62.4 %         | 83.9 %         | 6 944 584  | 3.88 MiB |
+
+    # Arguments
+    input_shape: Optional shape tuple, to be specified if you would like to use a model
+        with an input image resolution that is not (224, 224, 3).
+        It should have exactly 3 inputs channels.
+    input_tensor: optional Keras tensor (i.e. output of `layers.Input()`) to use as
+        image input for the model.
+    weights: one of `None` (random initialization), "imagenet" (pre-training on
+        ImageNet), or the path to the weights file to be loaded.
+    include_top: whether to include the fully-connected layer at the top of the network.
+    num_classes: optional number of classes to classify images into, only to be
+        specified if `include_top` is True, and if no `weights` argument is specified.
+
+    # Returns
+    A Keras model instance.
+
+    # Raises
+    ValueError: in case of invalid argument for `weights`, or invalid input shape.
+
+    # References
+    - [MeliusNet: Can Binary Neural Networks Achieve MobileNet-level Accuracy?](https://arxiv.org/abs/2001.05936)
+    """
+    return MeliusNet22Factory(
+        include_top=include_top,
+        weights=weights,
+        input_tensor=input_tensor,
+        input_shape=input_shape,
+        num_classes=num_classes,
+    ).build()
diff --git a/tests/models_test.py b/tests/models_test.py
@@ -37,6 +37,7 @@ def parametrize(func):
             (lqz.literature.BinaryDenseNet37, 640),
             (lqz.literature.BinaryDenseNet37Dilated, 640),
             (lqz.literature.BinaryDenseNet45, 800),
+            (lqz.literature.MeliusNet22, 512),
             (lqz.literature.XNORNet, 4096),
             (lqz.literature.DoReFaNet, 256),
             (lqz.literature.RealToBinaryNet, 512),
