My suspicion is that one quantizer is created and reused for the entire model, and flip_ratio looks at the same quantizer with inputs of different shapes and fails because of this.
from functools import partial
from typing import Callable, Tuple
import larq
import tensorflow as tf
from tensorflow import keras
from larq_zoo.sota import QuickNet
INPUT_SHAPE = 32
CLASSES_NUM = 10
EPOCHS = 100
BATCH_SIZE = 128
LEARNING_RATE = 5e-3
def quicknet(input_shape: int, num_classes: int) -> keras.Model:
quicknet_spatial_reduce_factor = 32
global_pool_shape = input_shape / quicknet_spatial_reduce_factor
quicknet_pretrained_base = QuickNet(input_shape=(input_shape, input_shape, 3), include_top=False)
quicknet_pretrained_base.trainable = True
return tf.keras.models.Sequential([
quicknet_pretrained_base,
keras.layers.AveragePooling2D(pool_size=(global_pool_shape, global_pool_shape)),
keras.layers.Flatten(),
keras.layers.Dense(num_classes, kernel_initializer="glorot_normal"),
tf.keras.layers.Activation("softmax", dtype="float32")
])
def get_dataset(batch_size: int, preprocessing: Callable) -> Tuple[tf.data.Dataset, tf.data.Dataset]:
train, test = tf.keras.datasets.cifar10.load_data()
train_dataset = (
tf.data.Dataset.from_tensor_slices(train).cache()
.shuffle(10 * batch_size, reshuffle_each_iteration=True)
.map(partial(preprocessing, training=True))
.batch(batch_size)
)
test_dataset = (
tf.data.Dataset.from_tensor_slices(test).cache()
.map(preprocessing)
.batch(batch_size)
)
return train_dataset, test_dataset
def identity_preprocess(x, y, training=False):
return x, y
if __name__ == "__main__":
with larq.context.metrics_scope(['flip_ratio']):
model = quicknet(INPUT_SHAPE, CLASSES_NUM)
larq.models.summary(model)
train_dataset, test_dataset = get_dataset(BATCH_SIZE, identity_preprocess)
optimizer = keras.optimizers.Adam(LEARNING_RATE)
loss = keras.losses.SparseCategoricalCrossentropy()
model.compile(
optimizer=optimizer, loss=loss, metrics=[keras.metrics.SparseCategoricalAccuracy()]
)
model.fit(
train_dataset, epochs=EPOCHS, validation_data=test_dataset
)
I would have expected the example to run without problems, as happens when quicknet is replaced with a series of binary operations, but instead I get the following error:
Traceback (most recent call last):
File "C:/Users/User/PycharmProjects/BNN-Playground/bug_replicate.py", line 59, in <module>
model = quicknet(INPUT_SHAPE, CLASSES_NUM)
File "C:/Users/User/PycharmProjects/BNN-Playground/bug_replicate.py", line 21, in quicknet
quicknet_pretrained_base = QuickNet(input_shape=(input_shape, input_shape, 3), include_top=False)
File "C:\Users\User\Anaconda3\lib\site-packages\larq_zoo\sota\quicknet.py", line 327, in QuickNet
num_classes=num_classes,
File "C:\Users\User\Anaconda3\lib\site-packages\zookeeper\core\factory.py", line 20, in wrapped_fn
result = fn(factory_instance)
File "C:\Users\User\Anaconda3\lib\site-packages\larq_zoo\sota\quicknet.py", line 189, in build
model = super().build()
File "C:\Users\User\Anaconda3\lib\site-packages\zookeeper\core\factory.py", line 20, in wrapped_fn
result = fn(factory_instance)
File "C:\Users\User\Anaconda3\lib\site-packages\larq_zoo\sota\quicknet.py", line 156, in build
x = self.residual_block(x, use_squeeze_and_excite)
File "C:\Users\User\Anaconda3\lib\site-packages\larq_zoo\sota\quicknet.py", line 104, in residual_block
x = self.conv_block(x, infilters, use_squeeze_and_excite)
File "C:\Users\User\Anaconda3\lib\site-packages\larq_zoo\sota\quicknet.py", line 92, in conv_block
)(x)
File "C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\keras\engine\base_layer.py", line 773, in __call__
outputs = call_fn(cast_inputs, *args, **kwargs)
File "C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\autograph\impl\api.py", line 237, in wrapper
raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:
C:\Users\User\Anaconda3\lib\site-packages\larq\layers_base.py:37 call *
return super().call(inputs)
C:\Users\User\Anaconda3\lib\site-packages\larq\layers_base.py:153 call *
return super().call(inputs)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\keras\layers\convolutional.py:209 call
outputs = self._convolution_op(inputs, self.kernel)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\nn_ops.py:1135 __call__
return self.conv_op(inp, filter)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\nn_ops.py:640 __call__
return self.call(inp, filter)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\nn_ops.py:239 __call__
name=self.name)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\nn_ops.py:2011 conv2d
name=name)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\gen_nn_ops.py:969 conv2d
data_format=data_format, dilations=dilations, name=name)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\framework\op_def_library.py:486 _apply_op_helper
(input_name, err))
ValueError: Tried to convert 'filter' to a tensor and failed. Error: in converted code:
C:\Users\User\Anaconda3\lib\site-packages\larq\quantizers.py:249 call *
return super().call(outputs)
C:\Users\User\Anaconda3\lib\site-packages\larq\quantizers.py:160 call *
self.add_metric(self.flip_ratio(inputs))
C:\Users\User\Anaconda3\lib\site-packages\larq\metrics.py:43 __call__ *
return super().__call__(inputs, **kwargs)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\keras\metrics.py:196 __call__
replica_local_fn, *args, **kwargs)
C:\Users\User\Anaconda3\lib\site-packages\larq\metrics.py:71 update_state *
unchanged_values = tf.math.count_nonzero(
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\util\dispatch.py:180 wrapper
return target(*args, **kwargs)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\math_ops.py:1305 equal
return gen_math_ops.equal(x, y, name=name)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\ops\gen_math_ops.py:3240 equal
name=name)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\framework\op_def_library.py:742 _apply_op_helper
attrs=attr_protos, op_def=op_def)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\framework\func_graph.py:595 _create_op_internal
compute_device)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\framework\ops.py:3322 _create_op_internal
op_def=op_def)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\framework\ops.py:1786 __init__
control_input_ops)
C:\Users\User\Anaconda3\lib\site-packages\tensorflow_core\python\framework\ops.py:1622 _create_c_op
raise ValueError(str(e))
ValueError: Dimensions must be equal, but are 64 and 128 for 'Equal' (op: 'Equal') with input shapes: [3,3,64,64], [3,3,128,128].