Handling Class Imbalance with AutoMM - Focal Loss¶
In this tutorial, we introduce how to use focal loss with the AutoMM package for balanced training. Focal loss is first introduced in this Paper and can be used for balancing hard/easy samples as well as un-even sample distribution among classes. This tutorial demonstrates how to use focal loss.
Create Dataset¶
We use the shopee dataset for demonstration in this tutorial. Shopee dataset contains 4 classes and has 200 samples each in the training set.
from autogluon.multimodal.utils.misc import shopee_dataset
download_dir = "./ag_automm_tutorial_imgcls_focalloss"
train_data, test_data = shopee_dataset(download_dir)
Downloading ./ag_automm_tutorial_imgcls_focalloss/file.zip from https://automl-mm-bench.s3.amazonaws.com/vision_datasets/shopee.zip...
/home/ci/autogluon/multimodal/src/autogluon/multimodal/data/templates.py:16: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
import pkg_resources
0%| | 0.00/84.0M [00:00<?, ?iB/s]
10%|▉ | 8.38M/84.0M [00:00<00:01, 43.0MiB/s]
20%|█▉ | 16.8M/84.0M [00:00<00:01, 56.2MiB/s]
30%|██▉ | 25.2M/84.0M [00:00<00:01, 52.4MiB/s]
40%|███▉ | 33.5M/84.0M [00:00<00:01, 45.6MiB/s]
48%|████▊ | 40.2M/84.0M [00:00<00:01, 41.5MiB/s]
53%|█████▎ | 44.5M/84.0M [00:01<00:00, 39.8MiB/s]
58%|█████▊ | 48.5M/84.0M [00:01<00:01, 34.4MiB/s]
62%|██████▏ | 52.0M/84.0M [00:01<00:00, 32.9MiB/s]
68%|██████▊ | 56.9M/84.0M [00:01<00:00, 35.8MiB/s]
72%|███████▏ | 60.6M/84.0M [00:01<00:00, 35.1MiB/s]
80%|███████▉ | 67.1M/84.0M [00:01<00:00, 41.2MiB/s]
90%|████████▉ | 75.5M/84.0M [00:01<00:00, 44.7MiB/s]
98%|█████████▊| 82.1M/84.0M [00:02<00:00, 38.6MiB/s]
100%|██████████| 84.0M/84.0M [00:02<00:00, 39.4MiB/s]
For the purpose of demonstrating the effectiveness of Focal Loss on imbalanced training data, we artificially downsampled the shopee training data to form an imbalanced distribution.
import numpy as np
import pandas as pd
ds = 1
imbalanced_train_data = []
for lb in range(4):
class_data = train_data[train_data.label == lb]
sample_index = np.random.choice(np.arange(len(class_data)), size=int(len(class_data) * ds), replace=False)
ds /= 3 # downsample 1/3 each time for each class
imbalanced_train_data.append(class_data.iloc[sample_index])
imbalanced_train_data = pd.concat(imbalanced_train_data)
print(imbalanced_train_data)
weights = []
for lb in range(4):
class_data = imbalanced_train_data[imbalanced_train_data.label == lb]
weights.append(1 / (class_data.shape[0] / imbalanced_train_data.shape[0]))
print(f"class {lb}: num samples {len(class_data)}")
weights = list(np.array(weights) / np.sum(weights))
print(weights)
image label
171 /home/ci/autogluon/docs/tutorials/multimodal/a... 0
60 /home/ci/autogluon/docs/tutorials/multimodal/a... 0
191 /home/ci/autogluon/docs/tutorials/multimodal/a... 0
109 /home/ci/autogluon/docs/tutorials/multimodal/a... 0
18 /home/ci/autogluon/docs/tutorials/multimodal/a... 0
.. ... ...
716 /home/ci/autogluon/docs/tutorials/multimodal/a... 3
695 /home/ci/autogluon/docs/tutorials/multimodal/a... 3
701 /home/ci/autogluon/docs/tutorials/multimodal/a... 3
620 /home/ci/autogluon/docs/tutorials/multimodal/a... 3
768 /home/ci/autogluon/docs/tutorials/multimodal/a... 3
[295 rows x 2 columns]
class 0: num samples 200
class 1: num samples 66
class 2: num samples 22
class 3: num samples 7
[np.float64(0.0239850482815907), np.float64(0.07268196448966878), np.float64(0.21804589346900635), np.float64(0.6852870937597342)]
Create and train MultiModalPredictor¶
Train with Focal Loss¶
We specify the model to use focal loss by setting the "optim.loss_func" to "focal_loss".
There are also three other optional parameters you can set.
optim.focal_loss.alpha - a list of floats which is the per-class loss weight that can be used to balance un-even sample distribution across classes.
Note that the len of the list must match the total number of classes in the training dataset. A good way to compute alpha for each class is to use the inverse of its percentage number of samples.
optim.focal_loss.gamma - float which controls how much to focus on the hard samples. Larger value means more focus on the hard samples.
optim.focal_loss.reduction - how to aggregate the loss value. Can only take "mean" or "sum" for now.
import uuid
from autogluon.multimodal import MultiModalPredictor
model_path = f"./tmp/{uuid.uuid4().hex}-automm_shopee_focal"
predictor = MultiModalPredictor(label="label", problem_type="multiclass", path=model_path)
predictor.fit(
hyperparameters={
"model.mmdet_image.checkpoint_name": "swin_tiny_patch4_window7_224",
"env.num_gpus": 1,
"optim.loss_func": "focal_loss",
"optim.focal_loss.alpha": weights, # shopee dataset has 4 classes.
"optim.focal_loss.gamma": 1.0,
"optim.focal_loss.reduction": "sum",
"optim.max_epochs": 10,
},
train_data=imbalanced_train_data,
)
predictor.evaluate(test_data, metrics=["acc"])
=================== System Info ===================
AutoGluon Version: 1.5.1b20260114
Python Version: 3.12.10
Operating System: Linux
Platform Machine: x86_64
Platform Version: #1 SMP Wed Mar 12 14:53:59 UTC 2025
CPU Count: 8
Pytorch Version: 2.9.1+cu128
CUDA Version: 12.8
GPU Memory: GPU 0: 14.57/14.57 GB
Total GPU Memory: Free: 14.57 GB, Allocated: 0.00 GB, Total: 14.57 GB
GPU Count: 1
Memory Avail: 28.39 GB / 30.95 GB (91.7%)
Disk Space Avail: 164.58 GB / 255.99 GB (64.3%)
===================================================
AutoMM starts to create your model. ✨✨✨
To track the learning progress, you can open a terminal and launch Tensorboard:
```shell
# Assume you have installed tensorboard
tensorboard --logdir /home/ci/autogluon/docs/tutorials/multimodal/advanced_topics/tmp/197fb5c3567f43bbbaf0ec690082376e-automm_shopee_focal
```
Seed set to 0
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[4], line 8
4 model_path = f"./tmp/{uuid.uuid4().hex}-automm_shopee_focal"
6 predictor = MultiModalPredictor(label="label", problem_type="multiclass", path=model_path)
----> 8 predictor.fit(
9 hyperparameters={
10 "model.mmdet_image.checkpoint_name": "swin_tiny_patch4_window7_224",
11 "env.num_gpus": 1,
12 "optim.loss_func": "focal_loss",
13 "optim.focal_loss.alpha": weights, # shopee dataset has 4 classes.
14 "optim.focal_loss.gamma": 1.0,
15 "optim.focal_loss.reduction": "sum",
16 "optim.max_epochs": 10,
17 },
18 train_data=imbalanced_train_data,
19 )
21 predictor.evaluate(test_data, metrics=["acc"])
File ~/autogluon/multimodal/src/autogluon/multimodal/predictor.py:540, in MultiModalPredictor.fit(self, train_data, presets, tuning_data, max_num_tuning_data, id_mappings, time_limit, save_path, hyperparameters, column_types, holdout_frac, teacher_predictor, seed, standalone, hyperparameter_tune_kwargs, clean_ckpts, predictions, labels, predictors)
537 assert isinstance(predictors, list)
538 learners = [ele if isinstance(ele, str) else ele._learner for ele in predictors]
--> 540 self._learner.fit(
541 train_data=train_data,
542 presets=presets,
543 tuning_data=tuning_data,
544 max_num_tuning_data=max_num_tuning_data,
545 time_limit=time_limit,
546 save_path=save_path,
547 hyperparameters=hyperparameters,
548 column_types=column_types,
549 holdout_frac=holdout_frac,
550 teacher_learner=teacher_learner,
551 seed=seed,
552 standalone=standalone,
553 hyperparameter_tune_kwargs=hyperparameter_tune_kwargs,
554 clean_ckpts=clean_ckpts,
555 id_mappings=id_mappings,
556 predictions=predictions,
557 labels=labels,
558 learners=learners,
559 )
561 return self
File ~/autogluon/multimodal/src/autogluon/multimodal/learners/base.py:665, in BaseLearner.fit(self, train_data, presets, tuning_data, time_limit, save_path, hyperparameters, column_types, holdout_frac, teacher_learner, seed, standalone, hyperparameter_tune_kwargs, clean_ckpts, **kwargs)
658 self.fit_sanity_check()
659 self.prepare_fit_args(
660 time_limit=time_limit,
661 seed=seed,
662 standalone=standalone,
663 clean_ckpts=clean_ckpts,
664 )
--> 665 fit_returns = self.execute_fit()
666 self.on_fit_end(
667 training_start=training_start,
668 strategy=fit_returns.get("strategy", None),
(...)
671 clean_ckpts=clean_ckpts,
672 )
674 return self
File ~/autogluon/multimodal/src/autogluon/multimodal/learners/base.py:577, in BaseLearner.execute_fit(self)
575 return dict()
576 else:
--> 577 attributes = self.fit_per_run(**self._fit_args)
578 self.update_attributes(**attributes) # only update attributes for non-HPO mode
579 return attributes
File ~/autogluon/multimodal/src/autogluon/multimodal/learners/base.py:1291, in BaseLearner.fit_per_run(self, max_time, save_path, ckpt_path, resume, enable_progress_bar, seed, hyperparameters, advanced_hyperparameters, config, df_preprocessor, data_processors, model, standalone, clean_ckpts)
1289 validation_metric, custom_metric_func = self.get_validation_metric_per_run()
1290 mixup_active, mixup_func = self.get_mixup_func_per_run(config=config)
-> 1291 loss_func, aug_loss_func = self.get_loss_func_per_run(config=config, mixup_active=mixup_active)
1292 model_postprocess_fn = self.get_model_postprocess_fn_per_run(loss_func=loss_func)
1293 num_gpus, strategy = self.get_num_gpus_and_strategy_per_run(config=config)
File ~/autogluon/multimodal/src/autogluon/multimodal/learners/base.py:845, in BaseLearner.get_loss_func_per_run(self, config, mixup_active)
844 def get_loss_func_per_run(self, config, mixup_active=None):
--> 845 loss_func = get_loss_func(
846 problem_type=self._problem_type,
847 mixup_active=mixup_active,
848 loss_func_name=config.optim.loss_func,
849 config=config.optim,
850 )
851 aug_loss_func = get_aug_loss_func(
852 config=config.optim,
853 problem_type=self._problem_type,
854 )
855 return loss_func, aug_loss_func
File ~/autogluon/multimodal/src/autogluon/multimodal/optim/losses/utils.py:63, in get_loss_func(problem_type, mixup_active, loss_func_name, config, **kwargs)
61 else:
62 if loss_func_name is not None and loss_func_name.lower() == "focal_loss":
---> 63 loss_func = FocalLoss(
64 alpha=config.focal_loss.alpha,
65 gamma=config.focal_loss.gamma,
66 reduction=config.focal_loss.reduction,
67 )
68 else:
69 loss_func = nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
File ~/autogluon/multimodal/src/autogluon/multimodal/optim/losses/focal_loss.py:43, in FocalLoss.__init__(self, alpha, gamma, reduction, eps)
40 self.reduction = reduction
41 self.eps = float(eps) if eps is not None else 1e-6
---> 43 alpha = self._parse_alpha(alpha)
44 self.nll_loss = nn.NLLLoss(weight=alpha, reduction="none")
File ~/autogluon/multimodal/src/autogluon/multimodal/optim/losses/focal_loss.py:62, in FocalLoss._parse_alpha(self, alpha)
59 except Exception:
60 raise ValueError(f"{type(alpha)} {alpha} is not in a supported format.")
---> 62 return torch.tensor(alpha, dtype=torch.float32)
ValueError: too many dimensions 'str'
Train without Focal Loss¶
import uuid
from autogluon.multimodal import MultiModalPredictor
model_path = f"./tmp/{uuid.uuid4().hex}-automm_shopee_non_focal"
predictor2 = MultiModalPredictor(label="label", problem_type="multiclass", path=model_path)
predictor2.fit(
hyperparameters={
"model.mmdet_image.checkpoint_name": "swin_tiny_patch4_window7_224",
"env.num_gpus": 1,
"optim.max_epochs": 10,
},
train_data=imbalanced_train_data,
)
predictor2.evaluate(test_data, metrics=["acc"])
As we can see that the model with focal loss is able to achieve a much better performance compared to the model without focal loss. When your data is imbalanced, try out focal loss to see if it brings improvements to the performance!
Citations¶
@misc{https://doi.org/10.48550/arxiv.1708.02002,
doi = {10.48550/ARXIV.1708.02002},
url = {https://arxiv.org/abs/1708.02002},
author = {Lin, Tsung-Yi and Goyal, Priya and Girshick, Ross and He, Kaiming and Dollár, Piotr},
keywords = {Computer Vision and Pattern Recognition (cs.CV), FOS: Computer and information sciences, FOS: Computer and information sciences},
title = {Focal Loss for Dense Object Detection},
publisher = {arXiv},
year = {2017},
copyright = {arXiv.org perpetual, non-exclusive license}
}