RT-DETR改进策略【小目标改进】| 2024-TOP 自适应阈值焦点损失(ATFL)提升对小目标的检测能力
一、本文介绍
本文记录的是
利用自适应阈值焦点损失(ATFL)优化RT-DETR目标检测的网络模型
。
ATFL
结合了
阈值解耦
的针对性和自适应调整损失权重的灵活性,
有效地处理目标与背景不平衡问题
。这一机制通过设置阈值
区分目标和背景样本
,并根据
预测概率值
自适应调整损失权重来
适应不同的样本特性
,提高了模型对目标特征的学习能力。在小目标检测中,
ATFL
可以被用于提升对小目标的检测能力,
特别是在目标和背景极度不平衡的情况下
。
二、自适应阈值焦点损失原理
自适应阈值焦点损失(ATFL)
是一种用于解决
红外小目标检测中目标与背景不平衡问题
的创新方法。以下将从其
设计出发点
、
原理
、
结构
和
优势
进行详细介绍,并结合文中相关图示辅助说明。
2.1 设计出发点
在红外图像中, 目标所占比例极小,背景占据主导 ,这使得模型在训练过程中更容易学习到背景特征,而忽略目标特征。传统的损失函数无法有效解决这种 样本不平衡问题 ,因此需要一种新的损失函数来促使模型 更加关注目标特征 。
2.2 原理
- 基于阈值的解耦 :通过 设置阈值 将易于识别的背景和难以识别的目标分离开。
例如,将预测概率值高于0.5的样本视为背景(容易样本),低于0.5的视为目标(困难样本)。
- 损失权重的调整 :对 不同类型的样本调整其损失权重 。对于目标样本,增强其损失,使模型更加关注目标特征;对于背景样本,减轻其损失。
- 自适应机制 :为了减少调整超参数带来的时间消耗,对超参数进行 自适应设计 。
例如,根据预测概率值和模型训练进度来动态调整焦点损失函数中的调制因子 γ 等参数。
2.3 结构
2.3.1 基础损失函数
从经典的
交叉熵损失函数
出发,其表达式为
L
B
C
E
=
−
(
y
l
o
g
(
p
)
+
(
1
−
y
)
l
o
g
(
1
−
p
)
)
\mathcal{L}_{BCE}=-(y log (p)+(1-y) log (1-p))
L
BCE
=
−
(
y
l
o
g
(
p
)
+
(
1
−
y
)
l
o
g
(
1
−
p
))
(其中
P
P
P
代表预测概率,
y
y
y
代表真实标签),也可简记为:
L
B
C
E
=
−
l
o
g
(
p
t
)
\mathcal{L}_{BCE}=-log \left(p_{t}\right)
L
BCE
=
−
l
o
g
(
p
t
)
其中
p
t
=
{
p
,
i
f
y
=
1
1
−
p
,
o
t
h
e
r
s
p_{t}= \begin{cases}p, & if y=1 \\ 1-p, & others \end{cases}
p
t
=
{
p
,
1
−
p
,
i
f
y
=
1
o
t
h
ers
2.3.2 焦点损失函数改进
焦点损失函数 F L ( p t ) = − ( 1 − p t ) γ l o g ( p t ) F L\left(p_{t}\right)=-\left(1-p_{t}\right)^{\gamma} log \left(p_{t}\right) F L ( p t ) = − ( 1 − p t ) γ l o g ( p t ) 通过引入调制因子 ( 1 − p t ) γ (1-p_{t})^{\gamma} ( 1 − p t ) γ 来降低易分类样本的损失贡献,但 在降低易样本损失的同时也会降低困难样本损失,不利于困难样本学习。
2.3.3 阈值焦点损失(TFL)函数
为解决上述问题,提出 T F L = { − ( λ − p t ) η l o g ( p t ) p t < = 0.5 − ( 1 − p t ) γ l o g ( p t ) p t > 0.5 T F L=\left\{\begin{array}{ll} -\left(\lambda-p_{t}\right)^{\eta} log \left(p_{t}\right) & p_{t}<=0.5 \\ -\left(1-p_{t}\right)^{\gamma} log \left(p_{t}\right) & p_{t}>0.5 \end{array}\right. TF L = { − ( λ − p t ) η l o g ( p t ) − ( 1 − p t ) γ l o g ( p t ) p t <= 0.5 p t > 0.5 其中 λ ( > 1 ) \lambda(>1) λ ( > 1 ) 为超参数, 对于不同数据集和模型需多次调整。
2.3.4 自适应阈值焦点损失(ATFL)函数
对 γ \gamma γ 和 η \eta η 进行 自适应调整 。 γ = − l n ( p ^ c ) \gamma=-ln \left(\hat{p}_{c}\right) γ = − l n ( p ^ c ) (其中 p ^ c = 0.05 × 1 t − 1 ∑ i = 0 t − 1 p i ‾ + 0.95 × p t \hat{p}_{c}=0.05 × \frac{1}{t-1} \sum_{i=0}^{t-1} \overline{p_{i}}+0.95 × p_{t} p ^ c = 0.05 × t − 1 1 ∑ i = 0 t − 1 p i + 0.95 × p t ), η = − l n ( p t ) \eta=-ln \left(p_{t}\right) η = − l n ( p t ) ,最终得到: A T F L = { − ( λ − p t ) − l n ( p t ) l o g ( p t ) p t < = 0.5 − ( 1 − p t ) − l n ( p ^ c ) l o g ( p t ) p t > 0.5 A T F L= \begin{cases}-\left(\lambda-p_{t}\right)^{-ln \left(p_{t}\right)} log \left(p_{t}\right) & p_{t}<=0.5 \\ -\left(1-p_{t}\right)^{-ln \left(\hat{p}_{c}\right)} log \left(p_{t}\right) & p_{t}>0.5\end{cases} A TF L = { − ( λ − p t ) − l n ( p t ) l o g ( p t ) − ( 1 − p t ) − l n ( p ^ c ) l o g ( p t ) p t <= 0.5 p t > 0.5
2.4 优势
- 缓解不平衡问题 :对于容易样本,随着 p t p_{t} p t 增加,损失值会按照设计要求降低,进一步减少了容易样本产生的损失,从而有效缓解了目标与背景的不平衡,使模型能够更好地学习目标特征。
- 自适应调整 :自适应机制避免了多次手动调整超参数,节省了时间成本,同时能够根据数据集和模型训练情况自动调整参数,提高了模型的性能和泛化能力。
- 提升困难样本学习 :通过增加困难样本(目标)的损失权重,使模型更加关注难以检测的目标,提高了对红外小目标的检测性能。
论文: https://arxiv.org/abs/2307.14723
源码: https://github.com/YangBo0411/infrared-small-target
三、自适应阈值焦点损失(ATFL)的实现代码
自适应阈值焦点损失(ATFL)
的实现代码如下:
class AdaptiveThresholdFocalLoss(nn.Module):
# Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
super(AdaptiveThresholdFocalLoss, self).__init__()
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
self.gamma = gamma
self.alpha = alpha
def forward(self, pred, true):
loss = self.loss_fcn(pred, true)
pred_prob = torch.sigmoid(pred)
p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
p_t = torch.Tensor(p_t)
mean_pt = p_t.mean()
p_t_list = []
p_t_list.append(mean_pt)
p_t_old = sum(p_t_list) / len(p_t_list)
p_t_new = 0.05 * p_t_old + 0.95 * mean_pt
# gamma =2
gamma = -torch.log(p_t_new)
p_t_high = torch.where(p_t > 0.5, (1.000001 - p_t) ** gamma, torch.zeros_like(p_t))
p_t_low = torch.where(p_t <= 0.5, (1.5 - p_t) ** (-torch.log(p_t)), torch.zeros_like(p_t)) # # 将两部分结果相加
modulating_factor = p_t_high + p_t_low
loss *= modulating_factor
return loss
四、添加步骤
4.1 修改ultralytics/utils/loss.py
此处需要修改的文件是
ultralytics/utils/loss.py
loss.py
中定义了模型的损失函数和计算方法,我们想要加入新的损失函数就只需要将代码放到这个文件内即可。
AdaptiveThresholdFocalLoss
添加后如下:
4.2 AdaptiveThresholdFocalLoss的调用
在
loss.py
中的
v8DetectionLoss函数
中添加如下代码,使模型调用此
AdaptiveThresholdFocalLoss函数
。
g = 1 # focal loss gamma
if g > 0:
self.bce = AdaptiveThresholdFocalLoss(self.bce, g)
此时再次训练模型便会使用
AdaptiveThresholdFocalLoss
计算模型的损失函数。
五、成功运行结果