##### Copyright 2020 Google LLC
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Codebook based quantization
Codebook based quantizaion is a non-uniform quantization technique that maps each weight or activation value to the index of a value in the codebook. This allows us to compress weights/activations even further with neglibible loss in performance. We will demonstrate this by training an object classification model and applying codebook quantization to the activation with the most values.
from keras.datasets import *
from keras.layers import *
from keras.models import Model
from keras.optimizers import *
from keras.regularizers import *
from keras.utils import to_categorical
from qkeras import *
from qkeras.codebook import *
def get_data(name, sample_size=1.0):
(x_train, y_train), (x_test, y_test) = globals()[name].load_data()
if len(x_train.shape) == 3:
x_train = x_train.reshape(x_train.shape + (1,))
x_test = x_test.reshape(x_test.shape + (1,))
x_train = x_train.astype("float32")
x_test = x_test.astype("float32")
mean = knp.mean(x_train, axis=(0, 1, 2, 3))
std = np.std(x_train, axis=(0, 1, 2, 3))
x_train = (x_train - mean) / (std + 1e-7)
x_test = (x_test - mean) / (std + 1e-7)
y_train_c = to_categorical(y_train, knp.max(y_train) + 1)
y_test_c = to_categorical(y_test, knp.max(y_test) + 1)
if sample_size != 1.0:
indexes = knp.asarray(range(x_train.shape[0]))
np.random.shuffle(indexes)
indexes = indexes[: int(x_train.shape[0] * sample_size)]
x_train = x_train[indexes]
y_train_c = y_train_c[indexes]
return (x_train, y_train_c), (x_test, y_test_c)
def get_model(
name,
X_train,
y_train,
X_test,
y_test,
blocks=[[32], [64], [128]],
quantizer_list=["quantized_relu_po2(4,4)", "quantized_relu_po2(4,4)"],
use_stochastic_rounding=0,
l1v=None,
epochs=10,
load_weights=True,
):
if l1v is None:
l1v = [0.0] * len(blocks)
X_shape = X_train.shape[1:]
x_i = x = Input(X_shape)
for b, block in enumerate(blocks):
# we are assuming we want to quantize the block that has sparsity
# so let's add dropout to the next layer
if b >= 1 and l1v[b - 1] != 0.0:
x = Dropout(0.3, name=f"drop{b}")(x)
for i in range(len(block)):
x = QConv2D(
block[i],
kernel_size=(3, 3),
strides=(2, 2),
padding="same",
kernel_quantizer=f"quantized_bits(4, use_stochastic_rounding={use_stochastic_rounding})",
bias_quantizer=f"quantized_po2(4, use_stochastic_rounding={use_stochastic_rounding})",
kernel_regularizer=l1(l1v[b]) if l1v[b] != 0.0 else None,
name=f"d{b}_{i}",
)(x)
if i != len(block) - 1:
if quantizer_list[b] in ["linear", "relu", "softmax", "sigmoid"]:
x = Activation(quantizer_list[b], name=f"a{b}_{i}")(x)
else:
x = QActivation(quantizer_list[b], name=f"a{b}_{i}")(x)
else:
x = QBatchNormalization(name=f"bn{b}_{i}")(x)
if b < len(blocks) - 1:
if quantizer_list[b] in ["linear", "relu", "softmax", "sigmoid"]:
x = Activation(quantizer_list[b], name=f"a{b}_{len(block)-1}")(x)
else:
x = QActivation(quantizer_list[b], name=f"a{b}_{len(block)-1}")(x)
else:
if len(block) > 0:
x = QActivation(
f"quantized_relu(6,2, use_stochastic_rounding={use_stochastic_rounding})",
name=f"a{b}_{len(block)-1}",
)(x)
x = Flatten(name="flatten")(x)
x = QDense(y_train.shape[1], name=f"d{len(blocks)-1}_{len(block)}")(x)
x = Activation("softmax", name=f"a{len(blocks)-1}_{len(block)}")(x)
model = Model(inputs=x_i, outputs=x)
model.summary()
model.compile(
loss="categorical_crossentropy", optimizer=Adam(0.001), metrics=["acc"]
)
try:
if load_weights and os.path.isfile(name + ".weights.h5"):
print("Found file...")
model.load_weights(name + ".weights.h5")
else:
model.fit(
X_train,
y_train,
validation_data=(X_test, y_test),
batch_size=128,
epochs=epochs,
verbose=2,
)
model.save_weights(name + ".weights.h5")
except:
model.fit(
X_train,
y_train,
validation_data=(X_test, y_test),
batch_size=128,
epochs=epochs,
verbose=2,
)
model.save_weights(name + ".weights.h5")
return model
name = "cifar10"
(X_train, y_train), (X_test, y_test) = get_data(name, sample_size=1)
model = get_model(
name,
X_train,
y_train,
X_test,
y_test,
blocks=[[32, 32], [64, 64], [128]],
quantizer_list=["quantized_relu(6,2)", "quantized_relu(6,2)"],
epochs=50,
load_weights=True,
)
/Users/mariuskoppel/cms/qkeras/venv/lib/python3.11/site-packages/keras/src/datasets/cifar.py:18: VisibleDeprecationWarning: dtype(): align should be passed as Python or NumPy boolean but got `align=0`. Did you mean to pass a tuple to create a subarray type? (Deprecated NumPy 2.4)
d = cPickle.load(f, encoding="bytes")
Model: "functional_1"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input_layer_1 (InputLayer) │ (None, 32, 32, 3) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d0_0 (QConv2D) │ (None, 16, 16, 32) │ 896 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a0_0 (QActivation) │ (None, 16, 16, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d0_1 (QConv2D) │ (None, 8, 8, 32) │ 9,248 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn0_1 (QBatchNormalization) │ (None, 8, 8, 32) │ 128 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a0_1 (QActivation) │ (None, 8, 8, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d1_0 (QConv2D) │ (None, 4, 4, 64) │ 18,496 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a1_0 (QActivation) │ (None, 4, 4, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d1_1 (QConv2D) │ (None, 2, 2, 64) │ 36,928 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn1_1 (QBatchNormalization) │ (None, 2, 2, 64) │ 256 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a1_1 (QActivation) │ (None, 2, 2, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d2_0 (QConv2D) │ (None, 1, 1, 128) │ 73,856 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn2_0 (QBatchNormalization) │ (None, 1, 1, 128) │ 512 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a2_0 (QActivation) │ (None, 1, 1, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ flatten (Flatten) │ (None, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d2_1 (QDense) │ (None, 10) │ 1,290 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a2_1 (Activation) │ (None, 10) │ 0 │ └─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 141,610 (553.16 KB)
Trainable params: 141,162 (551.41 KB)
Non-trainable params: 448 (1.75 KB)
Epoch 1/50
391/391 - 21s - 53ms/step - acc: 0.4149 - loss: 1.6186 - val_acc: 0.5019 - val_loss: 1.5106
Epoch 2/50
391/391 - 14s - 35ms/step - acc: 0.5516 - loss: 1.2503 - val_acc: 0.5745 - val_loss: 1.1772
Epoch 3/50
391/391 - 15s - 38ms/step - acc: 0.6058 - loss: 1.1104 - val_acc: 0.5640 - val_loss: 1.2066
Epoch 4/50
391/391 - 17s - 43ms/step - acc: 0.6320 - loss: 1.0632 - val_acc: 0.6137 - val_loss: 1.0699
Epoch 5/50
391/391 - 14s - 37ms/step - acc: 0.6511 - loss: 1.0010 - val_acc: 0.6250 - val_loss: 1.0728
Epoch 6/50
391/391 - 15s - 39ms/step - acc: 0.6819 - loss: 0.9200 - val_acc: 0.6406 - val_loss: 1.0122
Epoch 7/50
391/391 - 14s - 35ms/step - acc: 0.6984 - loss: 0.8537 - val_acc: 0.6528 - val_loss: 0.9860
Epoch 8/50
391/391 - 14s - 36ms/step - acc: 0.7144 - loss: 0.8085 - val_acc: 0.6550 - val_loss: 0.9842
Epoch 9/50
391/391 - 15s - 39ms/step - acc: 0.7297 - loss: 0.7673 - val_acc: 0.6596 - val_loss: 0.9843
Epoch 10/50
391/391 - 16s - 40ms/step - acc: 0.7410 - loss: 0.7346 - val_acc: 0.6619 - val_loss: 0.9936
Epoch 11/50
391/391 - 15s - 38ms/step - acc: 0.7494 - loss: 0.7109 - val_acc: 0.6439 - val_loss: 1.0735
Epoch 12/50
391/391 - 14s - 36ms/step - acc: 0.7498 - loss: 0.7181 - val_acc: 0.6647 - val_loss: 0.9700
Epoch 13/50
391/391 - 15s - 38ms/step - acc: 0.7614 - loss: 0.6803 - val_acc: 0.6648 - val_loss: 1.3618
Epoch 14/50
391/391 - 14s - 37ms/step - acc: 0.7703 - loss: 0.6620 - val_acc: 0.6752 - val_loss: 0.9417
Epoch 15/50
391/391 - 15s - 38ms/step - acc: 0.7812 - loss: 0.6208 - val_acc: 0.6713 - val_loss: 0.9850
Epoch 16/50
391/391 - 14s - 35ms/step - acc: 0.7922 - loss: 0.5931 - val_acc: 0.6670 - val_loss: 1.0179
Epoch 17/50
391/391 - 15s - 38ms/step - acc: 0.7958 - loss: 0.5746 - val_acc: 0.6575 - val_loss: 1.0581
Epoch 18/50
391/391 - 15s - 38ms/step - acc: 0.7995 - loss: 0.5625 - val_acc: 0.6724 - val_loss: 1.0153
Epoch 19/50
391/391 - 16s - 41ms/step - acc: 0.8094 - loss: 0.5431 - val_acc: 0.6827 - val_loss: 1.0291
Epoch 20/50
391/391 - 14s - 37ms/step - acc: 0.8141 - loss: 0.5209 - val_acc: 0.6703 - val_loss: 1.0595
Epoch 21/50
391/391 - 15s - 38ms/step - acc: 0.8091 - loss: 0.5362 - val_acc: 0.6629 - val_loss: 1.0729
Epoch 22/50
391/391 - 16s - 40ms/step - acc: 0.8094 - loss: 0.5468 - val_acc: 0.6493 - val_loss: 1.0504
Epoch 23/50
391/391 - 15s - 38ms/step - acc: 0.8256 - loss: 0.4942 - val_acc: 0.6809 - val_loss: 1.0269
Epoch 24/50
391/391 - 15s - 38ms/step - acc: 0.8392 - loss: 0.4527 - val_acc: 0.6773 - val_loss: 1.0432
Epoch 25/50
391/391 - 14s - 35ms/step - acc: 0.8409 - loss: 0.4460 - val_acc: 0.6830 - val_loss: 1.0521
Epoch 26/50
391/391 - 15s - 40ms/step - acc: 0.8488 - loss: 0.4242 - val_acc: 0.6591 - val_loss: 1.1478
Epoch 27/50
391/391 - 15s - 37ms/step - acc: 0.8507 - loss: 0.4141 - val_acc: 0.6726 - val_loss: 1.1206
Epoch 28/50
391/391 - 14s - 36ms/step - acc: 0.8554 - loss: 0.4029 - val_acc: 0.6678 - val_loss: 1.1668
Epoch 29/50
391/391 - 14s - 35ms/step - acc: 0.8546 - loss: 0.4028 - val_acc: 0.6661 - val_loss: 1.1881
Epoch 30/50
391/391 - 14s - 35ms/step - acc: 0.8550 - loss: 0.4005 - val_acc: 0.6719 - val_loss: 1.1821
Epoch 31/50
391/391 - 15s - 37ms/step - acc: 0.8478 - loss: 0.4274 - val_acc: 0.6514 - val_loss: 1.2850
Epoch 32/50
391/391 - 15s - 38ms/step - acc: 0.8514 - loss: 0.4242 - val_acc: 0.6570 - val_loss: 1.1238
Epoch 33/50
391/391 - 14s - 36ms/step - acc: 0.8513 - loss: 0.4165 - val_acc: 0.6790 - val_loss: 1.0588
Epoch 34/50
391/391 - 15s - 38ms/step - acc: 0.8556 - loss: 0.4081 - val_acc: 0.6718 - val_loss: 1.4007
Epoch 35/50
391/391 - 15s - 39ms/step - acc: 0.8581 - loss: 0.3971 - val_acc: 0.6759 - val_loss: 1.2020
Epoch 36/50
391/391 - 15s - 38ms/step - acc: 0.8658 - loss: 0.3688 - val_acc: 0.6657 - val_loss: 1.2170
Epoch 37/50
391/391 - 16s - 40ms/step - acc: 0.8710 - loss: 0.3572 - val_acc: 0.6780 - val_loss: 1.1904
Epoch 38/50
391/391 - 15s - 38ms/step - acc: 0.8741 - loss: 0.3481 - val_acc: 0.6699 - val_loss: 1.2108
Epoch 39/50
391/391 - 16s - 40ms/step - acc: 0.8781 - loss: 0.3338 - val_acc: 0.6754 - val_loss: 1.2237
Epoch 40/50
391/391 - 16s - 40ms/step - acc: 0.8792 - loss: 0.3332 - val_acc: 0.6676 - val_loss: 1.3164
Epoch 41/50
391/391 - 15s - 37ms/step - acc: 0.8776 - loss: 0.3339 - val_acc: 0.6652 - val_loss: 1.3028
Epoch 42/50
391/391 - 14s - 36ms/step - acc: 0.8748 - loss: 0.3383 - val_acc: 0.6768 - val_loss: 1.2733
Epoch 43/50
391/391 - 15s - 38ms/step - acc: 0.8824 - loss: 0.3202 - val_acc: 0.6625 - val_loss: 1.3783
Epoch 44/50
391/391 - 16s - 41ms/step - acc: 0.8782 - loss: 0.3337 - val_acc: 0.6585 - val_loss: 1.3520
Epoch 45/50
391/391 - 16s - 41ms/step - acc: 0.8702 - loss: 0.3674 - val_acc: 0.6738 - val_loss: 0.9334
Epoch 46/50
391/391 - 15s - 39ms/step - acc: 0.8786 - loss: 0.3411 - val_acc: 0.6769 - val_loss: 1.2422
Epoch 47/50
391/391 - 15s - 38ms/step - acc: 0.8814 - loss: 0.3291 - val_acc: 0.6736 - val_loss: 1.2969
Epoch 48/50
391/391 - 15s - 37ms/step - acc: 0.8904 - loss: 0.3028 - val_acc: 0.6762 - val_loss: 1.3155
Epoch 49/50
391/391 - 14s - 36ms/step - acc: 0.8906 - loss: 0.3020 - val_acc: 0.6779 - val_loss: 1.3146
Epoch 50/50
391/391 - 14s - 36ms/step - acc: 0.8922 - loss: 0.2931 - val_acc: 0.6670 - val_loss: 1.3430
from qkeras.codebook import *
cb_tables, models, km_models = activation_compression(
model,
{"loss": "categorical_crossentropy", "metrics": ["acc"]},
[2],
3,
X_train,
y_train,
X_test,
y_test,
sample_size=0.3,
)
Creating submodel...
Model: "functional_2"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input_layer_1 (InputLayer) │ (None, 32, 32, 3) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d0_0 (QConv2D) │ (None, 16, 16, 32) │ 896 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a0_0 (QActivation) │ (None, 16, 16, 32) │ 0 │ └─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 896 (3.50 KB)
Trainable params: 896 (3.50 KB)
Non-trainable params: 0 (0.00 B)
None
Creating submodel...
Model: "functional_3"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input_layer_2 (InputLayer) │ (None, 16, 16, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d0_1 (QConv2D) │ (None, 8, 8, 32) │ 9,248 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn0_1 (QBatchNormalization) │ (None, 8, 8, 32) │ 128 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a0_1 (QActivation) │ (None, 8, 8, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d1_0 (QConv2D) │ (None, 4, 4, 64) │ 18,496 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a1_0 (QActivation) │ (None, 4, 4, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d1_1 (QConv2D) │ (None, 2, 2, 64) │ 36,928 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn1_1 (QBatchNormalization) │ (None, 2, 2, 64) │ 256 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a1_1 (QActivation) │ (None, 2, 2, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d2_0 (QConv2D) │ (None, 1, 1, 128) │ 73,856 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn2_0 (QBatchNormalization) │ (None, 1, 1, 128) │ 512 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a2_0 (QActivation) │ (None, 1, 1, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ flatten (Flatten) │ (None, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ d2_1 (QDense) │ (None, 10) │ 1,290 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ a2_1 (Activation) │ (None, 10) │ 0 │ └─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 140,714 (549.66 KB)
Trainable params: 140,266 (547.91 KB)
Non-trainable params: 448 (1.75 KB)
None
sample_size: 0.3
fitting km[0]...
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 5s 3ms/step
313/313 ━━━━━━━━━━━━━━━━━━━━ 2s 5ms/step
Number of unique activations: 8
Evaluating...
313/313 - 5s - 17ms/step - acc: 0.6710 - loss: 1.3443
q = models[0].layers[-1].activation
in_table, out_table = cb_tables[0]
print(q)
print("in_table:", in_table)
print("out_table:", out_table)
quantized_relu(6,2)
in_table: [0. 0.1875 0.5 0.9375 1.4375 2.0625 2.875 3.8125]
out_table: [0 0 1 1 1 1 2 2 2 2 2 2 3 3 3 3 3 3 3 3 4 4 4 4 4 4 4 4 4 5 5 5 5 5 5 5 5
5 5 5 6 6 6 6 6 6 6 6 6 6 6 6 6 6 7 7 7 7 7 7 7 7 7 7]
for i, x in enumerate(q.range()):
print(f"{x:8}, {in_table[out_table[i]]:6}")
0.0, 0.0
0.0625, 0.0
0.125, 0.1875
0.1875, 0.1875
0.25, 0.1875
0.3125, 0.1875
0.375, 0.5
0.4375, 0.5
0.5, 0.5
0.5625, 0.5
0.625, 0.5
0.6875, 0.5
0.75, 0.9375
0.8125, 0.9375
0.875, 0.9375
0.9375, 0.9375
1.0, 0.9375
1.0625, 0.9375
1.125, 0.9375
1.1875, 0.9375
1.25, 1.4375
1.3125, 1.4375
1.375, 1.4375
1.4375, 1.4375
1.5, 1.4375
1.5625, 1.4375
1.625, 1.4375
1.6875, 1.4375
1.75, 1.4375
1.8125, 2.0625
1.875, 2.0625
1.9375, 2.0625
2.0, 2.0625
2.0625, 2.0625
2.125, 2.0625
2.1875, 2.0625
2.25, 2.0625
2.3125, 2.0625
2.375, 2.0625
2.4375, 2.0625
2.5, 2.875
2.5625, 2.875
2.625, 2.875
2.6875, 2.875
2.75, 2.875
2.8125, 2.875
2.875, 2.875
2.9375, 2.875
3.0, 2.875
3.0625, 2.875
3.125, 2.875
3.1875, 2.875
3.25, 2.875
3.3125, 2.875
3.375, 3.8125
3.4375, 3.8125
3.5, 3.8125
3.5625, 3.8125
3.625, 3.8125
3.6875, 3.8125
3.75, 3.8125
3.8125, 3.8125
3.875, 3.8125
3.9375, 3.8125
Weight compression using codebook quantization
conv_weights = model.layers[1].weights[0].numpy()
print(conv_weights.shape)
quantizer = model.layers[1].kernel_quantizer_internal
print(quantizer)
axis = 3
bits = 3
index_table, codebook_table = weight_compression(conv_weights, bits, axis, quantizer)
(3, 3, 3, 32)
quantized_bits(4,0,1,alpha='auto_po2')
32it [00:00, 146.06it/s]
print(codebook_table.shape)
codebook_table[0]
(32, 8)
array([-0.75 , -0.5 , -0.375, -0.125, 0.125, 0.25 , 0.375, 0.5 ])
print(index_table.shape)
index_table[:, :, :, 0]
(3, 3, 3, 32)
<tf.Tensor: shape=(3, 3, 3), dtype=int32, numpy=
array([[[2, 4, 1],
[3, 6, 4],
[4, 3, 3]],
[[0, 5, 4],
[3, 6, 5],
[5, 3, 3]],
[[5, 7, 7],
[3, 5, 3],
[4, 6, 3]]], dtype=int32)>
slices = []
for i in range(conv_weights.shape[axis]):
idx_slice = ops.take(index_table, i, axis=axis)
lut = ops.convert_to_tensor(codebook_table[i])
quantized_slice = ops.take(lut, idx_slice)
slices.append(quantized_slice)
new_conv_weights = ops.stack(slices, axis=axis)
new_conv_weights[:, :, :, 0]
<tf.Tensor: shape=(3, 3, 3), dtype=float64, numpy=
array([[[-0.375, 0.125, -0.5 ],
[-0.125, 0.375, 0.125],
[ 0.125, -0.125, -0.125]],
[[-0.75 , 0.25 , 0.125],
[-0.125, 0.375, 0.25 ],
[ 0.25 , -0.125, -0.125]],
[[ 0.25 , 0.5 , 0.5 ],
[-0.125, 0.25 , -0.125],
[ 0.125, 0.375, -0.125]]])>
conv_weights[:, :, :, 0]
array([[[-0.32802898, 0.10454272, -0.4944516 ],
[-0.2222067 , 0.35520843, 0.15265706],
[ 0.09113201, -0.1956394 , -0.06453487]],
[[-0.81179124, 0.20486663, 0.07929426],
[-0.2036663 , 0.3200098 , 0.26449645],
[ 0.28087974, -0.2313629 , -0.13119416]],
[[ 0.30737787, 0.49038434, 0.5390768 ],
[-0.18612136, 0.22579621, -0.14528836],
[ 0.10823875, 0.40017307, -0.19679691]]], dtype=float32)
bias = model.layers[1].weights[1].numpy()
model.layers[1].set_weights([new_conv_weights, bias])
model.evaluate(X_test, y_test)
313/313 ━━━━━━━━━━━━━━━━━━━━ 2s 5ms/step - acc: 0.6542 - loss: 1.4346
[1.4346091747283936, 0.65420001745224]