Repository: larq/larq

Issue:
Summary fails for subclassed keras models
### Describe the bug
When using subclassed Keras models, calling `larq.models.summary` fails and expects calling
`model.build` first.
Even if `build` is called, it fails with another less obvious error (seems to be a problem with output_shape)

### To Reproduce
```python
import larq
import tensorflow as tf
from tensorflow import keras
from larq_zoo.sota import QuickNet
from zookeeper import cli, task
from typing import Callable, Optional


class EmbedderWrapperModel(keras.Model):
    def __init__(self, zoo_class: Callable[..., keras.Model],
                 input_shape: int, num_classes: int, dynamic=False,
                 finetune_basenet=True, pretrained_basenet=True, cut_layer_name: Optional[str] = None):
        super(EmbedderWrapperModel, self).__init__(dynamic=dynamic)

        self.basenet = self._get_basenet(zoo_class, input_shape, finetune_basenet, pretrained_basenet, cut_layer_name)
        global_pool_shape = self.basenet.output_shape[1], self.basenet.output_shape[2]

        self.batch_norm = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
        self.global_pool = keras.layers.AveragePooling2D(pool_size=global_pool_shape)
        self.dense_softmax = keras.layers.Dense(num_classes, activation=tf.nn.softmax)

    def _get_basenet(self, zoo_class: Callable[..., keras.Model], input_shape: int,
                     finetune_basenet: bool, pretrained_basenet: bool, cut_layer_name: Optional[str]) -> keras.Model:
        weights = "imagenet" if pretrained_basenet else None

        if not cut_layer_name:
            basenet = zoo_class(input_shape=(input_shape, input_shape, 3), include_top=False, weights=weights)
        else:
            full_zoo_model = zoo_class(input_shape=(input_shape, input_shape, 3), include_top=True, weights=weights)
            inputs, outputs = full_zoo_model.inputs, full_zoo_model.get_layer(cut_layer_name).output
            basenet = keras.Model(inputs=inputs, outputs=outputs)

        basenet.trainable = finetune_basenet

        return basenet

    def call(self, inputs, training=False, mask=None):
        x = self.basenet(inputs, training=training)
        x = self.batch_norm(x, training=training)
        x = self.global_pool(x)
        x = keras.layers.Flatten()(x)
        x = self.dense_softmax(x)

        return x


@task
class BugTest:
    INPUT_SHAPE = 224
    CLASSES_NUM = 10

    FINETUNE_BASENET = False
    PRETRAINED_BASENET = True

    def run(self):
        cut_layer_name = "activation"  # `activation` is last relu layer before global pooling

        model = EmbedderWrapperModel(QuickNet, self.INPUT_SHAPE, self.CLASSES_NUM,
                                     finetune_basenet=self.FINETUNE_BASENET, pretrained_basenet=self.PRETRAINED_BASENET,
                                     cut_layer_name=cut_layer_name)

        model.build(input_shape=(None, 224, 224, 3))

        larq.models.summary(model)


if __name__ == "__main__":
    cli()
```

### Expected behavior
Expected `model.build` to solve the problem, instead got:
```
Traceback (most recent call last):
  File "C:/Users/User/PycharmProjects/BNN-Playground/summary_bug.py", line 68, in <module>
    cli()
  File "C:\Users\User\Anaconda3\lib\site-packages\click\core.py", line 829, in __call__
    return self.main(*args, **kwargs)
  File "C:\Users\User\Anaconda3\lib\site-packages\click\core.py", line 782, in main
    rv = self.invoke(ctx)
  File "C:\Users\User\Anaconda3\lib\site-packages\click\core.py", line 1259, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "C:\Users\User\Anaconda3\lib\site-packages\click\core.py", line 1066, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "C:\Users\User\Anaconda3\lib\site-packages\click\core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "C:\Users\User\Anaconda3\lib\site-packages\zookeeper\core\task.py", line 59, in command
    task_instance.run()
  File "C:/Users/User/PycharmProjects/BNN-Playground/summary_bug.py", line 64, in run
    larq.models.summary(model)
  File "C:\Users\User\Anaconda3\lib\site-packages\larq\models.py", line 466, in summary
    model_profile = ModelProfile(model)
  File "C:\Users\User\Anaconda3\lib\site-packages\larq\models.py", line 261, in __init__
    self.layer_profiles = [LayerProfile(l) for l in model.layers]
  File "C:\Users\User\Anaconda3\lib\site-packages\larq\models.py", line 261, in <listcomp>
    self.layer_profiles = [LayerProfile(l) for l in model.layers]
  File "C:\Users\User\Anaconda3\lib\site-packages\larq\models.py", line 158, in __init__
    n=p.count * self.output_pixels,
  File "C:\Users\User\Anaconda3\lib\site-packages\larq\models.py", line 227, in output_pixels
    if len(self.output_shape) == 4:
TypeError: object of type 'NoneType' has no len()
```

### Environment

TensorFlow version: 2.2.0rc3
Larq version: 0.9.4



Discussion:

Resolving patch:
diff --git a/larq/models.py b/larq/models.py
--- a/larq/models.py
+++ b/larq/models.py
@@ -150,7 +150,7 @@ def __init__(self, layer: tf.keras.layers.Layer):
 
         self.op_profiles = []
 
-        if isinstance(layer, mac_containing_layers):
+        if isinstance(layer, mac_containing_layers) and self.output_pixels:
             for p in self.weight_profiles:
                 if not p.is_bias():
                     self.op_profiles.append(
@@ -190,7 +190,10 @@ def op_count(
         if op_type != "mac":
             raise ValueError("Currently only counting of MAC-operations is supported.")
 
-        if isinstance(self._layer, op_count_supported_layer_types):
+        if (
+            isinstance(self._layer, op_count_supported_layer_types)
+            and self.output_pixels
+        ):
             count = 0
             for op in self.op_profiles:
                 if (precision is None or op.precision == precision) and (
@@ -222,14 +225,15 @@ def output_shape_str(self) -> str:
             return "?"
 
     @property
-    def output_pixels(self) -> int:
+    def output_pixels(self) -> Optional[int]:
         """Number of pixels for a single feature map (1 for fully connected layers)."""
+        if not self.output_shape:
+            return None
         if len(self.output_shape) == 4:
             return int(np.prod(self.output_shape[1:3]))
-        elif len(self.output_shape) == 2:
+        if len(self.output_shape) == 2:
             return 1
-        else:
-            raise NotImplementedError()
+        raise NotImplementedError()
 
     @property
     def unique_param_bidtwidths(self) -> Sequence[int]:


Test patch:
diff --git a/larq/models_test.py b/larq/models_test.py
--- a/larq/models_test.py
+++ b/larq/models_test.py
@@ -6,6 +6,23 @@
 from larq.models import ModelProfile
 
 
+class ToyModel(tf.keras.Model):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.conv = lq.layers.QuantConv2D(
+            filters=32,
+            kernel_size=(3, 3),
+            kernel_quantizer="ste_sign",
+            input_shape=(64, 64, 1),
+            padding="same",
+        )
+        self.pool = tf.keras.layers.GlobalAvgPool2D()
+        self.dense = tf.keras.layers.Dense(10, activation="softmax")
+
+    def call(self, inputs):
+        return self.dense(self.pool(self.conv(inputs)))
+
+
 def get_profile_model():
     return tf.keras.models.Sequential(
         [
@@ -129,6 +146,14 @@ def test_summary(snapshot, capsys):
     snapshot.assert_match(captured.out)
 
 
+def test_subclass_model_summary(snapshot, capsys):
+    model = ToyModel()
+    model.build((None, 32, 32, 3))
+    lq.models.summary(model)
+    captured = capsys.readouterr()
+    snapshot.assert_match(captured.out)
+
+
 def test_summary_invalid_model():
     with pytest.raises(ValueError):
         lq.models.summary(tf.keras.Model())
diff --git a/larq/snapshots/snap_models_test.py b/larq/snapshots/snap_models_test.py
--- a/larq/snapshots/snap_models_test.py
+++ b/larq/snapshots/snap_models_test.py
@@ -54,3 +54,25 @@
 | Number of MACs                 0      |
 +---------------------------------------+
 '''
+
+snapshots['test_subclass_model_summary 1'] = '''+toy_model stats-------------------------------------------------------------+
+| Layer                     Input prec.   Outputs  # 1-bit  # 32-bit  Memory |
+|                                 (bit)                x 1       x 1    (kB) |
++----------------------------------------------------------------------------+
+| quant_conv2d                        -  multiple      864        32    0.23 |
+| global_average_pooling2d            -  multiple        0         0       0 |
+| dense                               -  multiple        0       330    1.29 |
++----------------------------------------------------------------------------+
+| Total                                                864       362    1.52 |
++----------------------------------------------------------------------------+
++toy_model summary------------------------+
+| Total params                   1.23 k   |
+| Trainable params               1.23 k   |
+| Non-trainable params           0        |
+| Model size                     1.52 KiB |
+| Model size (8-bit FP weights)  470.00 B |
+| Float-32 Equivalent            4.79 KiB |
+| Compression Ratio of Memory    0.32     |
+| Number of MACs                 0        |
++-----------------------------------------+
+'''