MS COCO 2017
(categories: 20, train: 28474, val: 1219)v1
, integrating advanced features of nn.Module
into the IouLoss
.Fire-Smoke
(categories: 2, train: 4320, val: 1330)After initializing your model, initialize the IouLoss
and assign it to the model. The instance attribute iou_mean
of IouLoss
will be output to the state_dict
for saving and loading during the training process.
# After initializing your model, integrate the IoU loss module into the model.
m = YourDetectionModel()
m.iouloss = IouLoss(ltype='WIoU', monotonous=False)
# Ensure that `IouLoss` is assigned to the model before executing `load_state_dict`.
m.load_state_dict(torch.load('last.pt'))
Modify the bounding box regression loss in the loss function. Note that the confidence of the bounding box uses $IoU$ instead of $1-\mathcal{L}_{WIoU}$.
m.iouloss.train()
iloss, liou = m.iouloss(xywh2xyxy(pred), xywh2xyxy(gt), ret_iou=True)
# Accumulate the regression loss.
lbox += iloss.mean()
# Adjust the objectness term to use IoU as the bounding box confidence.
tobj = 1 - liou.detach()
@article{tong2023wise,
title={Wise-IoU: Bounding Box Regression Loss with Dynamic Focusing Mechanism},
author={Tong, Zanjia and Chen, Yuhang and Xu, Zewei and Yu, Rong},
journal={arXiv preprint arXiv:2301.10051},
year={2023}
}