AutoQKeras define goal

[LordScarface]

Hello and thank you for your work of the Project!

I'm trying to quantize the Encoder Stage of a Convolutional Autoencoder with AutoQKeras. My goal is to reduce the number of bits in those Layers. Now I'm not sure how to correctly call AutoQKeras in order to minimize the loss and get the best Model. My current Code:

import keras
from keras import layers
import tensorflow as tf

from qkeras.autoqkeras import *
from qkeras import *
from qkeras.utils import model_quantize
from qkeras.qtools import run_qtools
from qkeras.qtools import settings as qtools_settings

import tempfile

[...] # load the dataset

IMG_SHAPE = datagen[0][0][0].shape # shape = (96, 96, 1)

z_dim = 16

# Build Autoencoder
autoencoder = keras.Sequential(name='Autoencoder')

autoencoder.add(Input(shape=IMG_SHAPE))

autoencoder.add(Conv2D(32, 4, strides=2, activation='relu', padding='same'))
autoencoder.add(Conv2D(64, 4, strides=2, activation='relu', padding='same'))
autoencoder.add(Conv2D(128, 4, strides=2, activation='relu', padding='same'))
autoencoder.add(Conv2D(256, 4, strides=2, activation='relu', padding='same'))
autoencoder.add(Dense(z_dim))

autoencoder.add(Dense(z_dim))
autoencoder.add(Conv2DTranspose(128, 5, strides=2, activation='relu', padding='same'))
autoencoder.add(Conv2DTranspose(64, 5, strides=2, activation='relu', padding='same'))
autoencoder.add(Conv2DTranspose(32, 6, strides=2, activation='relu', padding='same'))
autoencoder.add(Conv2DTranspose(1, 6, strides=2, activation='sigmoid', padding='same'))

autoencoder.compile(optimizer='adam', loss='binary_crossentropy')

physical_devices = tf.config.list_physical_devices()
for d in physical_devices:
    print(d)
    

has_tpus = np.any([d.device_type == "TPU" for d in physical_devices])

if has_tpus:
    TPU_WORKER = 'local'

    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
        tpu=TPU_WORKER, job_name='tpu_worker')
    if TPU_WORKER != 'local':
        tf.config.experimental_connect_to_cluster(resolver, protocol='grpc+loas')
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.experimental.TPUStrategy(resolver)
    print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

    cur_strategy = strategy
else:
    cur_strategy = tf.distribute.get_strategy()
    
custom_objects = {}

quantization_config = {
        "kernel": {
                "binary": 1,
                "stochastic_binary": 1,
                "ternary": 2,
                "stochastic_ternary": 2,
                "quantized_bits(2,1,1,alpha=1.0)": 2,
                "quantized_bits(4,0,1,alpha=1.0)": 4,
                "quantized_bits(8,0,1,alpha=1.0)": 8
        },
        "bias": {
                "quantized_bits(4,0,1)": 4,
                "quantized_bits(8,3,1)": 8
        },
        "activation": {
                "binary": 1,
                "ternary": 2,
                "quantized_relu(3,1)": 3,
                "quantized_relu(4,2)": 4,
                "quantized_relu(8,2)": 8,
                "quantized_relu(8,4)": 8,
                "quantized_relu(16,8)": 16
        },
        "linear": {
                "binary": 1,
                "ternary": 2,
                "quantized_bits(4,1)": 4,
                "quantized_bits(8,2)": 8,
                "quantized_bits(16,10)": 16
        }
}

limit = {
    "Dense": [8, 8, 4],
    "Conv2D": [4, 8, 4],
    "DepthwiseConv2D": [4, 8, 4],
    "Activation": [4]
}

run_config = {
  "output_dir": tempfile.mkdtemp(),
  "quantization_config": quantization_config,
  "learning_rate_optimizer": False,
  "transfer_weights": False,
  "mode": "random",
  "seed": 42,
  "limit": limit,
  "tune_filters": "layer",
  "tune_filters_exceptions": "^dense",
  "distribution_strategy": cur_strategy,
  # first layer is input, layer two layers are softmax and flatten
  "layer_indexes": range(0, 4),
  "max_trials": 20
}

print("quantizing layers:", [autoencoder.layers[i].name for i in run_config["layer_indexes"]])

autoqk = AutoQKeras(autoencoder, custom_objects=custom_objects, **run_config)
autoqk.fit(x_data, x_data, validation_data=(x_data_val, x_data_val), batch_size=1, epochs=20)

qmodel = autoqk.get_best_model()

optimizer = Adam(learning_rate=0.02)
qmodel.compile(optimizer=optimizer, loss="binary_crossentropy")
qmodel.fit(x_data, x_data, epochs=20, batch_size=1, validation_data=(x_data_val, x_data_val))

score = qmodel.evaluate(x_data_val, x_data_val, verbose=1)
print(score)

When I train the Autoencoder on the Dataset without using AutoQKeras it works fine, however after quantizing and retrieving the best Model, the output of the prediction is all black.

I suspect I need to pass some argument, that AutoQKeras knows it should minimize the loss?

Best Regards, Lukas

[danielemoro]

Hi! I think you are missing the goal parameter in your run_config dictionary. Please see the tutorial notebook for AutoQKeras and search for the key term "goal" to see an example of how it can be used https://github.com/google/qkeras/blob/master/notebook/AutoQKeras.ipynb

AutoQKeras will then try to optimize that goal when searching for quantization strategies.