How to use tensorflow (tflite) model maker with class weights

I am using the tflite model maker for Object Detection and I would really like to apply class weights to my training process to fight class imbalance. I am using the Efficient-Det-Lite model family.

However, I do not know how to do this or if this is even possible.

I checked the github source code of the hparams config and I cannot find anything saying class weights here.

I tried setting the class_weight parameter directly in the object_detector_spec.py like this:

model.fit(
    train_dataset,
    epochs=epochs,
    class_weight = class_weights
    steps_per_epoch=steps_per_epoch,
    callbacks=train_lib.get_callbacks(config.as_dict(), val_dataset),
    validation_data=val_dataset,
    validation_steps=validation_steps)

but the output is the error

ValueError: `class_weight` is only supported for Models with a single output.

So my question is, it is possible with minor effort to apply class weights to training using TF Lite Model Maker? If yes, how can I do this?

My current code for training looks like this

import os

from tflite_model_maker.config import QuantizationConfig
from tflite_model_maker.config import ExportFormat
from tflite_model_maker import model_spec
from tflite_model_maker import object_detector

import tensorflow as tf
assert tf.__version__.startswith('2')

tf.get_logger().setLevel('ERROR')
from absl import logging
logging.set_verbosity(logging.ERROR)

model_name = 'efficientdet-lite0'
custom_model_dir_name = 'model_'+"1907"
epochs = 50
batch_size = 16
model_dir = 'models/'+model_name+'/'+custom_model_dir_name+'_e'+str(epochs)+'_b'+str(batch_size)
spec = object_detector.EfficientDetLite0Spec(
    model_name = model_name,
    model_dir='/home/alex/checkpoints/',
    hparams='grad_checkpoint=true,strategy=gpus',
    epochs=epochs, batch_size=batch_size,
    steps_per_execution=1, moving_average_decay=0,
    var_freeze_expr='(efficientnet|fpn_cells|resample_p6)',
    tflite_max_detections=25
)

train_data, validation_data, test_data = object_detector.DataLoader.from_csv(file_path)

model = object_detector.create(train_data, model_spec=spec, train_whole_model=True, validation_data=validation_data)

Solution

My answer:

No there is no easy way. The issue ValueError: `class_weight` is only supported for Models with a single output. is a problem that is still in discussion in an open issue in the tensorflow repository.

One way to solve is to severly downgrade tensorflow and do not wish to got that far. However it is of course possible to modify the loss function applied to the model yourself and apply weights to them to get the desired result.

What I did in the end was using the pre-defined focal-loss in my model which accounts for class imbalance which parameters γ and α. The α-balanced-focal-loss is a good way to handle imbalance in a dataset for instance in a case where not enough training samples are there.