Quantization aware training - label smoothing drops performance?
Hi @mhyeonsoo 

Could you please share a reproducible code in order to better understand and investigate the issue?

Thanks. @pjpratik 
Sure, here's my all code for training.

```python
def build_model(args):
    IMG_SHAPE = (args.input_dim, args.input_dim, 3)
    # Transfer learning model with MobileNetV3
    base_model = tf.keras.applications.MobileNetV3Large(
        input_shape=IMG_SHAPE,
        include_top=False,
        weights='imagenet',
        minimalistic=True,
#         include_preprocessing=False
    )
    # Freeze the pre-trained model weights
    base_model.trainable = False
    x = tf.keras.layers.GlobalMaxPooling2D()(base_model.output)
    x = tf.keras.layers.Dropout(0.2, name="top_dropout")(x)
    x = tf.keras.layers.Dense(args.num_classes, activation='softmax')(x)
    
    model = tf.keras.Model(base_model.input, x)

    return model


def setup_pretrained_model(args, ckpt_path):
    """
    Function to load pretrained model
    """
    
    model = build_model(args)
    model.load_weights(ckpt_path)
    
    return model


def run_model(args, ckpt_path):
    # build base model architecture
    model = build_model(args)
    label_list = os.listdir(args.input_path)
    
    # Compile the model    
    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        args.lr,
        decay_steps=100000,
        decay_rate=0.96,
        staircase=True)
    model.compile(optimizer=tf.keras.optimizers.Lion(lr=args.lr), 
                  loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=args.label_smoothing),
                  metrics=['accuracy']
                  )
    
    # load data
    n_sample, train_ds, val_ds = load_data(args, args.input_dim)

    
    hist = model.fit(train_ds,
                 epochs=args.epochs,
                 validation_data=val_ds,
                 steps_per_epoch=n_sample // args.batch_size,
                 validation_steps=val_ds.n // args.batch_size,
                 callbacks=[checkpoint_callback(ckpt_path),
                            early_stopping(),
                            logging_callback(args.log_dir),
                            save_metadata_callback(ckpt_path, label_list)],
                 verbose=1)


def fit_model(args):
    save_model_file = args.model_path + '/model_best.h5'
    qat_model_file = args.model_path + '/qat_model_best.h5'
    label_list = os.listdir(args.input_path)
    # fine-tunining baseline model for few epochs
    history = run_model(args, save_model_file)
    
    # load trained weights for quantization aware training
    model = setup_pretrained_model(args, save_model_file)
    
    def apply_quantization_to_dense(layer):
        if isinstance(layer, tf.keras.layers.Dense):
            return tfmot.quantization.keras.quantize_annotate_layer(layer)
        return layer

    annotated_model = tf.keras.models.clone_model(
        model,
        clone_function=apply_quantization_to_dense,
    )
    
    # Now that the Dense layers are annotated,
    # `quantize_apply` actually makes the model quantization aware.
    quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)

    n_sample, train_ds, val_ds = load_data(args, args.input_dim)

    # `quantize_model` requires a recompile.
    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        args.lr / 10.0,
        decay_steps=100000,
        decay_rate=0.96,
        staircase=True)
    
    # Compile the model
    quant_aware_model.compile(optimizer=tf.keras.optimizers.Lion(lr=lr_schedule), 
                  loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=args.label_smoothing),
                  metrics=['accuracy']
                  )
    
    hist = quant_aware_model.fit(train_ds,
                                 epochs=args.epochs,
                                 validation_data=val_ds,
                                 steps_per_epoch=n_sample // args.batch_size,
                                 validation_steps=val_ds.n // args.batch_size,
                                 callbacks=[checkpoint_callback(qat_model_file),
                                            early_stopping(),
                                            logging_callback(args.log_dir),
                                            save_metadata_callback(qat_model_file, label_list)],
                                 verbose=1)

    plot_hist(hist)
    quant_aware_model.save(args.model_path + '/qat_model_last.h5')
    
``` Hi @mhyeonsoo, do you have logs showing the performance differences? If your dataset is small enough and shareable that'll be easiest. Also can you provide the code used to call the above functions that show the performance difference? Thanks for your help! @pkgoogle 
sorry for the late response.
I am using private data so I am not allowed to share the dataset.
But yes, I am using small amount of data (~50 to 100 for each class).

I think as for the log, I can share this evaluation results. (includes confidence score / prediction result)
1. For the non-Label-smoothing trained model
![image](https://github.com/tensorflow/tensorflow/assets/32179857/c4926794-9e69-4115-bcc6-b59f08ad0fe9)

2. For the Label-smoothing trained model
![image](https://github.com/tensorflow/tensorflow/assets/32179857/cb044aa5-1b7d-4946-a32a-3ae571ca9632)


It seemed like the model does not work at all for few data which shows 0.125 confidence (I am using 8 classes total) Hi @mhyeonsoo,

I am having trouble reproducing your issue, where are you getting your callbacks? (checkpoint_callback, early_stopping, logging_callback, and save_metadata_callback), how is load_data defined? Here is a slightly modified version of your code, if you can help me resolve these definitions or add more context we can investigate better. Here is a [gist](https://colab.sandbox.google.com/gist/pkgoogle/e7d48935d27c869b5578a1144b7ae650/61783.ipynb#scrollTo=5fGNAN4Znm61) that starts that process, if you can return to me a working version which shows your issue, that'll be easiest. I used MNIST here to get us started, it seems like the issue might be data independent but this will help us figure that out. Thanks for your help. Hi @pkgoogle,

Here is the callback functions that I used.

```python
class SaveMetadataCallback(Callback):
    def __init__(self, filepath, labels):
        super(SaveMetadataCallback, self).__init__()
        self.filepath = filepath
        self.labels = labels

    def on_epoch_end(self, epoch, logs=None):
        # Save after each epoch
        with h5py.File(self.filepath, 'a') as f:
            if 'label_names' not in f:
                label_dataset = f.create_dataset('label_names', data=self.labels)

    def on_train_end(self, logs=None):
        # OR Save after training completion
        with h5py.File(self.filepath, 'a') as f:
            if 'label_names' not in f:
                label_dataset = f.create_dataset('label_names', data=self.labels)


def save_metadata_callback(ckpt_path, labels):
    return SaveMetadataCallback(filepath=ckpt_path, labels=labels)

# Create a callback that saves the model's checkpoints
def checkpoint_callback(ckpt_path):
    return tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_path,
                                              monitor='val_loss',
                                              verbose=1,
                                              save_best_only=True,
                                              save_weights_only=False,
                                              mode='auto',
                                              save_freq='epoch')



def early_stopping():
    return tf.keras.callbacks.EarlyStopping(monitor='val_loss',
                                            patience=10,
                                            verbose=1)


def logging_callback(log_dir):
    return tf.keras.callbacks.TensorBoard(log_dir=log_dir,
                                          histogram_freq=0,
                                          write_graph=True,
                                          write_images=False,
                                          update_freq='epoch',
                                          profile_batch=2,
                                          embeddings_freq=0,
                                          embeddings_metadata=None)
```


And for the dataloader, I followed tensorflow keras ImageDataGenerator.
```python
def load_data(args, input_size):
    train_datagen = ImageDataGenerator(
        rotation_range=20,
        zoom_range=0.05,
        width_shift_range=0.05,
        height_shift_range=0.05,
        shear_range=0.05,
        horizontal_flip=True,
        fill_mode="nearest",
        brightness_range=[0.8, 1.2],
        preprocessing_function=tray_crop,
        validation_split=0.2)
```


Please tell me if there is anything else needed.

Thanks Hi @mhyeonsoo,

How do you define the class "Callback"? is ImageDataGenerator tf.keras.preprocessing.image.ImageDataGenerator or something else? where/how is tray_crop defined?

Alternatively, if you can resolve the issues in this gist we can reduce the amount of turn around: [gist](https://colab.sandbox.google.com/gist/pkgoogle/b4e011c99e76665d79e75a8b10ef0d3c/61783.ipynb)  @pkgoogle 
My apologies for the incomplete codes. I needed some filtering due to the private information.
I setup the codes in the gist, so it may work without an error! The callback class was just a subclassing Class of keras Callback. It seems like it had not been saved well. I saved the [gist](https://colab.research.google.com/gist/mhyeonsoo/3b0d0f05128f893aaca467f24ef0b656/61783.ipynb) again. Please check the updates. Thanks! Hi @mhyeonsoo, your load_data function doesn't return anything and you try to unpack 3 values (n_sample, train_ds, val_ds) from it... did you finish it but not include the return statement? Thanks for your help. Hi @pkgoogle, I used custom directory_based image dataloader so it returned successfully, but since we reproduce with the MNIST here, I think we can just use ds_train and ds_test instead.
I modified the [gist](https://colab.research.google.com/gist/mhyeonsoo/3b0d0f05128f893aaca467f24ef0b656/61783.ipynb), so please test with this one.
 Hi @mhyeonsoo, I made a decent amount of changes to try to reproduce this: [gist](https://colab.sandbox.google.com/gist/pkgoogle/7c9af594527e8c4bc227dd9573d3e2e5/61783.ipynb)

I got the label_list from ds_train directly, I am unsure if this is what you meant by label_list

It seems like the model explicitly expects a batch dimension. I am running into the current issue you can see on the gist (can't batch inconsistent image sizes). I switched to imagenette as that's closer to how mobilenetv3 is trained. Can you give me some statistics about your dataset to see how to best progress on this? (i.e. what are the sizes of your dataset) Is it consistent? Any other information you think will be useful? Hi @pkgoogle, sorry for these delays. It must be weigh easier to reproduce if I could share the whole dataset, my apologies.
As for the label_list, I just removed the part that is using label_list because it is not mandatory.
[gist](https://colab.research.google.com/gist/mhyeonsoo/5747d34914914cdd34831ebe81e0e4f7/61783.ipynb)
For the image info, I resized the image as 224x224. And the whole dimension of it is (1,224,224,3) which is 3 channel.
Since the number of my dataset is about ~700, I used 16 as batch size.

I am not sure how tfds.load works here, but for mine, since I directly used ImageDataGenerator format, there had been a target_size argument so that I could resize it with this.

Again, sorry for making this inconvenience, and please tell me if anything is needed.

Thanks Hi @mhyeonsoo, I was able to get things through this way: [gist](https://colab.sandbox.google.com/gist/pkgoogle/24c1bd9a459174143e051ce2ca1f7817/61783.ipynb)

Lower label_smoothing in general performs better. Imagenette's labeling is likely not incorrect, so I wouldn't expect label_smoothing to help. If you think your labeling is more likely to be incorrect then it may help (at least on validation accuracy). Of course if your labeling is more incorrect than 10% then I recommend improving your labeling instead. So it essentially doesn't make sense to have label_smoothing > 10% (0.1) anyways.

are you using label_smoothing = 1.0? or near 1.0? If you are then it makes sense that your model can't tell the classes apart, which is why P(c) = 1/num_classes. This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. Hi @pkgoogle 
Thanks for testing all these things. I was wondering about the possibility of the performance fluctuation depending on the dataset or baseline model architecture.

As you mentioned, it may be because of my dataset (it's quite small amount, pretty clear to understand, etc). 

Again, thanks for considering responses and reproducues for this issue.

Thanks! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61783">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61783">No</a>

### 1. System information

- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): ubuntu18.04
- TensorFlow installation (pip package or built from source): pip package
- TensorFlow library (version, if pip package or github SHA, if built from source): 2.13.0

### 2. Code
```python
# fine-tunining baseline model for few epochs
history = run_model(args, save_model_file)
    
# load trained weights for quantization aware training
model = setup_pretrained_model(args, save_model_file)
    
def apply_quantization_to_dense(layer):
    if isinstance(layer, tf.keras.layers.Dense):
        return tfmot.quantization.keras.quantize_annotate_layer(layer)
    return layer

annotated_model = tf.keras.models.clone_model(
    model,
    clone_function=apply_quantization_to_dense,
)

# Build Model
annotated_model.build((None, args.input_dim, args.input_dim ,3))

# Now that the Dense layers are annotated,
# `quantize_apply` actually makes the model quantization aware.
quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)

#     quant_aware_model.summary()
n_sample, train_ds, val_ds = load_data(args, args.input_dim)

# `quantize_model` requires a recompile.
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    args.lr / 10.0,
    decay_steps=100000,
    decay_rate=0.96,
    staircase=True)
    
# Compile the model
quant_aware_model.compile(optimizer=tf.keras.optimizers.Lion(lr=lr_schedule), 
              loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=args.label_smoothing),
              metrics=['accuracy']
              )
```

### 3. Failure after conversion
model successfully converted into tflite with the example code flow of quantization aware training,
but the test accuracy for tflite model is so low (~80%) compared to .pb model file. (~97%)

Before adopting label smoothing, quantization aware trained model (tflite) was working as good as original .pb model.

Is there any relationship between label smoothing and quantization aware training process?

Thanks!
