# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf 
import tensorflow_hub as tfhub
import tensorflow_datasets as tfds

# +
(raw_train, raw_validation, raw_test), metadata = tfds.load(
    'cats_vs_dogs',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
)

def format_image(image, label):
    image = tf.image.resize(image, (224, 224)) / 255.0
    return  image, label

BATCH_SIZE = 16
num_examples = metadata.splits['train'].num_examples
num_classes = metadata.features['label'].num_classes

train_batches = (raw_train.shuffle(num_examples // 4)
                 .map(format_image).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE))

validation_batches = (raw_validation
                      .map(format_image)
                      .batch(BATCH_SIZE)
                      .prefetch(tf.data.AUTOTUNE))

test_batches = raw_test.map(format_image).batch(BATCH_SIZE)

# -

# ### Model

from tensorflow.keras.applications.resnet50 import ResNet50 

# +
base_model = ResNet50(input_shape=(224,224,3),
                         include_top=False,pooling='avg')

# Customize the output layers
flatten_all = tf.keras.layers.Flatten()(base_model.output)
Dense_1 = tf.keras.layers.Dense(units=512,activation='relu')(flatten_all)
prediction_layer = tf.keras.layers.Dense(units=2, activation='softmax')(Dense_1)

# Concatenate the model
model = tf.keras.models.Model(inputs=base_model.input, outputs=prediction_layer)

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

# hist = model.fit(train_batches,
#                  epochs=EPOCHS,
#                  validation_data=validation_batches)

# model.summary()
# -

# ## 

# ## Quantization Aware Training with Keras Models

import tensorflow_model_optimization as tfmot

tfmot.quantization.keras.quantiz

# +
import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

# q_aware stands for for quantization aware.
q_aware_model = quantize_model(model)

# `quantize_model` requires a recompile.
q_aware_model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

# q_aware_model.summary()
q_aware_model.fit(train_batches, validation_data=validation_batches, epochs=5,)
# -

tf.saved_model.save(q_aware_model,'Test/')

# +
converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

quantized_tflite_model = converter.convert()
# -

with open('test.tflite','wb') as f:
    f.write(quantized_tflite_model)



StarcoderdataJupyterScriptsDedupFiltered