qkeras.utils.clone_model is broken?

I have a problem cloning a model. The code bellow shows an example.

import tensorflow as tf
import numpy as np
import qkeras
import matplotlib.pyplot as plt 

from tensorflow.keras.datasets import mnist

# Setup train and test splits
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# Flatten the images
image_vector_size = 28*28
num_classes = 10  # ten unique digits
x_train = x_train.reshape(x_train.shape[0], image_vector_size)
x_train = x_train.astype('float32')
x_test = x_test.reshape(x_test.shape[0], image_vector_size)
x_test = x_test.astype('float32')

y_train = tf.one_hot(y_train, 10)
y_test = tf.one_hot(y_test, 10)

model = tf.keras.models.Sequential()
# We don't loose any info here since mnist are 8-bit gray-scale images. we just add this quantization
# to explicitly encode this for the chisel4ml optimizer.
model.add(tf.keras.layers.Input(shape=image_vector_size))
model.add(qkeras.QActivation(qkeras.quantized_bits(bits=8, integer=8, keep_negative=False, alpha="auto_po2")))

model.add(qkeras.QDense(32, kernel_quantizer=qkeras.quantized_bits(bits=4, integer=3, keep_negative=True, alpha="auto_po2")))
model.add(tf.keras.layers.BatchNormalization())
model.add(qkeras.QActivation(qkeras.quantized_relu(bits=4, integer=3)))

model.add(qkeras.QDense(32, kernel_quantizer=qkeras.quantized_bits(bits=4, integer=3, keep_negative=True, alpha="auto_po2")))
model.add(tf.keras.layers.BatchNormalization())
model.add(qkeras.QActivation(qkeras.quantized_relu(bits=4, integer=3)))

model.add(qkeras.QDense(32, kernel_quantizer=qkeras.quantized_bits(bits=4, integer=3, keep_negative=True, alpha="auto_po2")))
model.add(tf.keras.layers.BatchNormalization())
model.add(qkeras.QActivation(qkeras.quantized_relu(bits=4, integer=3)))

model.add(qkeras.QDense(num_classes, kernel_quantizer=qkeras.quantized_bits(bits=4, integer=3, keep_negative=True, alpha="auto_po2"), activation='softmax'))

model.compile(optimizer="adam",
              loss='categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, batch_size=32, epochs=1, verbose=True)
nmodel = qkeras.utils.clone_model(model)

nmodel.layers[1].kernel_quantizer

The output of this is:

{'class_name': 'quantized_bits',
 'config': {'bits': 4,
  'integer': 3,
  'symmetric': True,
  'alpha': 'auto_po2',
  'keep_negative': True,
  'use_stochastic_rounding': False,
  'qnoise_factor': 1.0}}

So what is actually returned for the field "kernel_quantizer" is a dictionary. Additionally, the field "kernel" is non-existent in nmodel.layers[1]. Am I using this function right? or is there some other function I could use?

(some background) What I am actually trying to do is modify the layers, so create a new model using different layers of the original model. I first tried just

nlayers = copy.deepcopy(model.layers)
nlayers = modify(nlayers)
new_model = tf.keras.models.Sequential(nlayers)

But this gave me an error (from the deepcopy function) cannot pickle '_thread.RLock' object.

Does anyone know a better way of doing this? Thanks and regards, Jure

google / qkeras

qkeras.utils.clone_model is broken? #95