Balanced MSE
Code for the paper:
Balanced MSE for Imbalanced Visual Regression
Jiawei Ren, Mingyuan Zhang, Cunjun Yu, Ziwei Liu
CVPR 2022 (Oral)
News
- [07/28/2022] We have released the code of IHMR in MMHuman3d
- [04/27/2022] We have released the code for the Synthetic Benchmark, including multi-dimensional Balanced MSE and visualizations!
- [03/31/2022] Code release, with an interactive demo and a hands-on tutorial.
Live Demo
Check out our live demo in the Hugging Face
Tutorial
We provide a minimal working example of Balanced MSE using the BMC implementation on a small-scale dataset, Boston Housing dataset.
The notebook is developed on top of Deep Imbalanced Regression (DIR) Tutorial, we thank the authors for their amazing tutorial!
Quick Preview
A code snippet of the Balanced MSE loss is shown below. We use the BMC implementation for demonstration, BMC does not require any label prior beforehand.
One-dimensional Balanced MSE
def bmc_loss(pred, target, noise_var):
"""Compute the Balanced MSE Loss (BMC) between `pred` and the ground truth `targets`.
Args:
pred: A float tensor of size [batch, 1].
target: A float tensor of size [batch, 1].
noise_var: A float number or tensor.
Returns:
loss: A float tensor. Balanced MSE Loss.
"""
logits = - (pred - target.T).pow(2) / (2 * noise_var) # logit size: [batch, batch]
loss = F.cross_entropy(logits, torch.arange(pred.shape[0])) # contrastive-like loss
loss = loss * (2 * noise_var).detach() # optional: restore the loss scale, 'detach' when noise is learnable
return lossnoise_var is a one-dimensional hyper-parameter. noise_var can be optionally optimized in training:
class BMCLoss(_Loss):
def __init__(self, init_noise_sigma):
super(BMCLoss, self).__init__()
self.noise_sigma = torch.nn.Parameter(torch.tensor(init_noise_sigma))
def forward(self, pred, target):
noise_var = self.noise_sigma ** 2
return bmc_loss(pred, target, noise_var)
criterion = BMCLoss(init_noise_sigma)
optimizer.add_param_group({'params': criterion.noise_sigma, 'lr': sigma_lr, 'name': 'noise_sigma'})Multi-dimensional Balanced MSE
The multi-dimensional implementation is compatible with the 1-D version.
from torch.distributions import MultivariateNormal as MVN
def bmc_loss_md(pred, target, noise_var):
"""Compute the Multidimensional Balanced MSE Loss (BMC) between `pred` and the ground truth `targets`.
Args:
pred: A float tensor of size [batch, d].
target: A float tensor of size [batch, d].
noise_var: A float number or tensor.
Returns:
loss: A float tensor. Balanced MSE Loss.
"""
I = torch.eye(pred.shape[-1])
logits = MVN(pred.unsqueeze(1), noise_var*I).log_prob(target.unsqueeze(0)) # logit size: [batch, batch]
loss = F.cross_entropy(logits, torch.arange(pred.shape[0])) # contrastive-like loss
loss = loss * (2 * noise_var).detach() # optional: restore the loss scale, 'detach' when noise is learnable
return lossnoise_var is still a one-dimensional hyper-parameter and can be optionally learned in training.
Run Experiments
Please go into the sub-folder to run experiments.
As for IHMR, we have released our code and pretrained models in MMHuman3d
Citation
@inproceedings{ren2021bmse,
title={Balanced MSE for Imbalanced Visual Regression},
author={Ren, Jiawei and Zhang, Mingyuan and Yu, Cunjun and Liu, Ziwei},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
year={2022}
}Acknowledgment
This work is supported by NTU NAP, MOE AcRF Tier 2 (T2EP20221-0033), the National Research Foundation, Singapore under its AI Singapore Programme, and under the RIE2020 Industry Alignment Fund β Industry Collabo- ration Projects (IAF-ICP) Funding Initiative, as well as cash and in-kind contribution from the industry partner(s).
The code is developed on top of Delving into Deep Imbalanced Regression.