def prepare_qat_model(model_in_path, model_out_path, train_images, train_labels):
    import tensorflow_model_optimization as tfmot
    quantize_model = tfmot.quantization.keras.quantize_model
    model = tf.keras.models.load_model(model_in_path)
    q_aware_model = quantize_model(model)
    q_aware_model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
    train_images_subset = train_images[0:1000]
    train_labels_subset = train_labels[0:1000]
    q_aware_model.fit(train_images_subset, train_labels_subset, batch_size=500, epochs=1, validation_split=0.1)
    q_aware_model.save(model_out_path)
def prepare_qat_model(model_in_path, model_out_path, train_images, train_labels):
    import tensorflow_model_optimization as tfmot
    quantize_model = tfmot.quantization.keras.quantize_model
    model = tf.keras.models.load_model(model_in_path)
    q_aware_model = quantize_model(model)
    q_aware_model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
    train_images_subset = train_images[0:1000]
    train_labels_subset = train_labels[0:1000]
    q_aware_model.fit(train_images_subset, train_labels_subset, batch_size=500, epochs=1, validation_split=0.1)
    q_aware_model.save(model_out_path)