Don't set quantizers as fields

Author: lgeiger

larq_zoo/core/model_factory.py

@@ -1,26 +1,17 @@
-from typing import Callable, Optional, Tuple, Union
+from typing import Optional, Tuple
 
 import tensorflow as tf
 from zookeeper import ComponentField, Field
 from zookeeper.tf import Dataset
 
 from larq_zoo.core import utils
 
-QuantizerType = Union[
-    tf.keras.layers.Layer, Callable[[tf.Tensor], tf.Tensor], str, None
-]
-ConstraintType = Union[tf.keras.constraints.Constraint, str, None]
 DimType = Optional[int]
 
 
 class ModelFactory:
     """A base class for Larq Zoo models. Defines some common fields."""
 
-    # Don't set any defaults here.
-    input_quantizer: QuantizerType = Field()
-    kernel_quantizer: QuantizerType = Field()
-    kernel_constraint: ConstraintType = Field()
-
     # This field is included for automatic inference of `num_clases`, if no
     # value is otherwise provided. We set `allow_missing` because we don't want
     # to throw an error if a dataset is not provided, as long as `num_classes`

larq_zoo/literature/binary_alex_net.py

@@ -17,9 +17,9 @@ class BinaryAlexNetFactory(ModelFactory):
 
     inflation_ratio: int = Field(1)
 
-    input_quantizer = Field("ste_sign")
-    kernel_quantizer = Field("ste_sign")
-    kernel_constraint = Field("weight_clip")
+    input_quantizer = "ste_sign"
+    kernel_quantizer = "ste_sign"
+    kernel_constraint = "weight_clip"
 
     def conv_block(
         self,

larq_zoo/literature/birealnet.py

@@ -14,9 +14,9 @@ class BiRealNetFactory(ModelFactory):
 
     filters: int = Field(64)
 
-    input_quantizer = Field("approx_sign")
-    kernel_quantizer = Field("magnitude_aware_sign")
-    kernel_constraint = Field("weight_clip")
+    input_quantizer = "approx_sign"
+    kernel_quantizer = "magnitude_aware_sign"
+    kernel_constraint = "weight_clip"
 
     kernel_initializer: Union[tf.keras.initializers.Initializer, str] = Field(
         "glorot_normal"

larq_zoo/literature/densenet.py

@@ -21,9 +21,17 @@ class BinaryDenseNet(tf.keras.models.Model):
 class BinaryDenseNetFactory(ModelFactory):
     """Implementation of [BinaryDenseNet](https://arxiv.org/abs/1906.08637)"""
 
-    input_quantizer = Field(lambda: lq.quantizers.SteSign(clip_value=1.3))
-    kernel_quantizer = Field(lambda: lq.quantizers.SteSign(clip_value=1.3))
-    kernel_constraint = Field(lambda: lq.constraints.WeightClip(clip_value=1.3))
+    @property
+    def input_quantizer(self):
+        return lq.quantizers.SteSign(clip_value=1.3)
+
+    @property
+    def kernel_quantizer(self):
+        return lq.quantizers.SteSign(clip_value=1.3)
+
+    @property
+    def kernel_constraint(self):
+        return lq.constraints.WeightClip(clip_value=1.3)
 
     initial_filters: int = Field(64)
     growth_rate: int = Field(64)

larq_zoo/literature/dorefanet.py

@@ -38,11 +38,15 @@ class DoReFaNetFactory(ModelFactory):
 
     activations_k_bit: int = Field(2)
 
-    input_quantizer = Field(
-        lambda self: lq.quantizers.DoReFaQuantizer(k_bit=self.activations_k_bit)
-    )
-    kernel_quantizer = Field(lambda: magnitude_aware_sign_unclipped)
-    kernel_constraint = Field(None)
+    @property
+    def input_quantizer(self):
+        return lq.quantizers.DoReFaQuantizer(k_bit=self.activations_k_bit)
+
+    @property
+    def kernel_quantizer(self):
+        return magnitude_aware_sign_unclipped
+
+    kernel_constraint = None
 
     def conv_block(
         self, x, filters, kernel_size, strides=1, pool=False, pool_padding="same"

larq_zoo/literature/resnet_e.py

@@ -15,9 +15,17 @@ class BinaryResNetE18Factory(ModelFactory):
     num_layers: int = Field(18)
     initial_filters: int = Field(64)
 
-    input_quantizer = Field(lambda: lq.quantizers.SteSign(clip_value=1.25))
-    kernel_quantizer = Field(lambda: lq.quantizers.SteSign(clip_value=1.25))
-    kernel_constraint = Field(lambda: lq.constraints.WeightClip(clip_value=1.25))
+    @property
+    def input_quantizer(self):
+        return lq.quantizers.SteSign(clip_value=1.25)
+
+    @property
+    def kernel_quantizer(self):
+        return lq.quantizers.SteSign(clip_value=1.25)
+
+    @property
+    def kernel_constraint(self):
+        return lq.constraints.WeightClip(clip_value=1.25)
 
     @property
     def spec(self):

larq_zoo/literature/xnornet.py

@@ -24,9 +24,9 @@ def xnor_weight_scale(x):
 class XNORNetFactory(ModelFactory):
     """Implementation of [XNOR-Net](https://arxiv.org/abs/1603.05279)"""
 
-    input_quantizer = Field("ste_sign")
-    kernel_quantizer = Field("xnor_weight_scale")
-    kernel_constraint = Field("weight_clip")
+    input_quantizer = "ste_sign"
+    kernel_quantizer = "xnor_weight_scale"
+    kernel_constraint = "weight_clip"
 
     kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = Field(
         lambda: tf.keras.regularizers.l2(5e-7)

larq_zoo/sota/quicknet.py

@@ -68,9 +68,17 @@ class QuickNetBaseFactory(ModelFactory):
     transition_block: MethodType = Field(None)
     stem_filters: int = Field(64)
 
-    input_quantizer = Field(lambda: lq.quantizers.SteSign(clip_value=1.25))
-    kernel_quantizer = Field(lambda: lq.quantizers.SteSign(clip_value=1.25))
-    kernel_constraint = Field(lambda: lq.constraints.WeightClip(clip_value=1.25))
+    @property
+    def input_quantizer(self):
+        return lq.quantizers.SteSign(clip_value=1.25)
+
+    @property
+    def kernel_quantizer(self):
+        return lq.quantizers.SteSign(clip_value=1.25)
+
+    @property
+    def kernel_constraint(self):
+        return lq.constraints.WeightClip(clip_value=1.25)
 
     def conv_block(
         self, x: tf.Tensor, filters: int, use_squeeze_and_excite: bool, strides: int = 1