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/meliusnet.py

@@ -26,9 +26,18 @@ class MeliusNetFactory(ModelFactory):
     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))
+
+    @property
+    def input_quantizer(self):
+        return lq.quantizers.SteSign(1.3)
+
+    @property
+    def kernel_quantizer(self):
+        return lq.quantizers.SteSign(1.3)
+
+    @property
+    def kernel_constraint(self):
+        return 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)

larq_zoo/literature/real_to_bin_nets.py

@@ -11,7 +11,7 @@
 from zookeeper import Field, factory
 
 from larq_zoo.core import utils
-from larq_zoo.core.model_factory import ModelFactory, QuantizerType
+from larq_zoo.core.model_factory import ModelFactory
 
 
 class _SharedBaseFactory(ModelFactory, metaclass=ABCMeta):
@@ -126,8 +126,8 @@ class StrongBaselineNetFactory(_SharedBaseFactory):
 
     scaling_r: int = 8
 
-    input_quantizer: QuantizerType = Field(None)
-    kernel_quantizer: QuantizerType = Field(None)
+    input_quantizer = None
+    kernel_quantizer = None
 
     class LearnedRescaleLayer(tf.keras.layers.Layer):
         """Implements the learned activation rescaling XNOR-Net++ style.
@@ -359,53 +359,68 @@ def block(
 @factory
 class StrongBaselineNetBANFactory(StrongBaselineNetFactory):
     model_name = Field("baseline_ban")
-    input_quantizer = Field("ste_sign")
-    kernel_quantizer = Field(None)
-    kernel_constraint = Field(None)
-    kernel_regularizer = Field(lambda: tf.keras.regularizers.l2(1e-5))
+    input_quantizer = "ste_sign"
+    kernel_quantizer = None
+    kernel_constraint = None
+
+    @property
+    def kernel_regularizer(self):
+        return tf.keras.regularizers.l2(1e-5)
 
 
 @factory
 class StrongBaselineNetBNNFactory(StrongBaselineNetFactory):
     model_name = Field("baseline_bnn")
-    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"
 
 
 @factory
 class RealToBinNetFPFactory(RealToBinNetFactory):
     model_name = Field("r2b_fp")
-    input_quantizer = Field(lambda: tf.keras.layers.Activation("tanh"))
-    kernel_quantizer = Field(None)
-    kernel_constraint = Field(None)
-    kernel_regularizer = Field(lambda: tf.keras.regularizers.l2(1e-5))
+    kernel_quantizer = None
+    kernel_constraint = None
+
+    @property
+    def input_quantizer(self):
+        return tf.keras.layers.Activation("tanh")
+
+    @property
+    def kernel_regularizer(self):
+        return tf.keras.regularizers.l2(1e-5)
 
 
 @factory
 class RealToBinNetBANFactory(RealToBinNetFactory):
     model_name = Field("r2b_ban")
-    input_quantizer = Field("ste_sign")
-    kernel_quantizer = Field(None)
-    kernel_constraint = Field(None)
-    kernel_regularizer = Field(lambda: tf.keras.regularizers.l2(1e-5))
+    input_quantizer = "ste_sign"
+    kernel_quantizer = None
+    kernel_constraint = None
+
+    @property
+    def kernel_regularizer(self):
+        return tf.keras.regularizers.l2(1e-5)
 
 
 @factory
 class RealToBinNetBNNFactory(RealToBinNetFactory):
     model_name = Field("r2b_bnn")
-    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"
 
 
 @factory
 class ResNet18FPFactory(ResNet18Factory):
     model_name = Field("resnet_fp")
-    input_quantizer = Field(None)
-    kernel_quantizer = Field(None)
-    kernel_constraint = Field(None)
-    kernel_regularizer = Field(lambda: tf.keras.regularizers.l2(1e-5))
+    input_quantizer = None
+    kernel_quantizer = None
+    kernel_constraint = None
+
+    @property
+    def kernel_regularizer(self):
+        return tf.keras.regularizers.l2(1e-5)
 
 
 def RealToBinaryNet(

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

@@ -65,9 +65,17 @@ class QuickNetBaseFactory(ModelFactory, abc.ABC):
     transition_block: Callable[..., tf.Tensor] = Field()
     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 __post_configure__(self):
         assert (

larq_zoo/training/knowledge_distillation/knowledge_distillation.py

@@ -4,7 +4,7 @@
 import tensorflow as tf
 from zookeeper import ComponentField, Field, factory
 
-from larq_zoo.core.model_factory import ConstraintType, ModelFactory, QuantizerType
+from larq_zoo.core.model_factory import ModelFactory
 
 
 class AttentionMatchingLossLayer(tf.keras.layers.Layer):
@@ -296,9 +296,9 @@ class TeacherStudentModelFactory(ModelFactory):
     teacher_model: tf.keras.models.Model = ComponentField(allow_missing=True)
     student_model: tf.keras.models.Model = ComponentField()
 
-    input_quantizer: QuantizerType = Field(allow_missing=True)
-    kernel_quantizer: QuantizerType = Field(allow_missing=True)
-    kernel_constraint: ConstraintType = Field(allow_missing=True)
+    input_quantizer = None
+    kernel_quantizer = None
+    kernel_constraint = None
 
     # Must be set if there is a teacher and allow_missing teacher weights is not True.
     # Either a full path or the name of a network (in which case it will be sought in the current `model_dir`).