YOLO26 自定义损失函数 重写 init_criterion 方法 损失类不继承基类

YOLO26 自定义损失函数 重写 init_criterion 方法 损失类不继承基类

flyfish

YOLO 的 ClassificationModel 基类中，专门设计了 init_criterion 方法：模型初始化时，会自动调用该方法生成损失函数并赋值给 self.criterion 属性。

通过继承模型基类，重写 init_criterion 方法，让模型在初始化阶段就原生创建自定义损失函数。

二、实现步骤

步骤1：自定义损失类

步骤2：自定义模型类，重写损失初始化方法

继承 ClassificationModel 基类，重写 init_criterion 方法，返回自定义损失的实例。
模型初始化时会自动调用该方法，将自定义损失设为模型的原生属性。

步骤3：自定义训练器，重写模型创建方法

继承官方 ClassificationTrainer，重写 get_model 方法，返回自定义的模型类实例，让训练器全程使用自定义模型进行训练。

步骤4：配置参数，启动训练

按照官方格式配置超参数，使用自定义训练器启动训练。

代码

import torch
from torch.nn import CrossEntropyLoss
from ultralytics.models.yolo.classify.train import ClassificationTrainer
from ultralytics.nn.tasks import ClassificationModel

class FocalClassificationLoss:
"""
自定义聚焦损失函数（Focal Loss）
完全兼容YOLO分类框架接口
"""
def __init__(self, gamma=2.0, alpha=None, label_smoothing=0.1):
"""
初始化损失函数超参数
:param gamma: 聚焦因子，调节难分样本的关注权重
:param alpha: 类别平衡权重
:param label_smoothing: 标签平滑系数
"""
self.gamma = gamma
self.alpha = torch.tensor(alpha) if alpha else None
self.label_smoothing = label_smoothing
self.ce = CrossEntropyLoss(label_smoothing=label_smoothing, reduction="none")
self.call_count = 0

def __call__(self, preds, batch):
"""
YOLO框架标准损失计算接口
:param preds: 模型预测输出
:param batch: 批次数据，包含标签
:return: 反向传播用损失、日志用损失
"""
self.call_count += 1

# 兼容多输出格式，提取有效预测张量
preds = preds[1] if isinstance(preds, (list, tuple)) else preds
# 计算基础交叉熵损失
ce_loss = self.ce(preds, batch["cls"])
# 计算真实标签的预测概率
pt = torch.exp(–ce_loss)
# Focal Loss计算
focal_loss = ((1 – pt) ** self.gamma) * ce_loss

# 应用类别平衡权重
if self.alpha is not None:
alpha_t = self.alpha.to(preds.device)[batch["cls"]]
focal_loss = alpha_t * focal_loss

loss = focal_loss.mean()
return loss, loss.detach()

class CustomClassificationModel(ClassificationModel):
"""
自定义分类模型
继承官方基类，重写损失函数初始化方法，原生支持自定义损失
"""
def init_criterion(self):
"""
重写官方损失初始化方法
模型初始化时会自动调用该方法创建损失函数
"""
return FocalClassificationLoss(gamma=2.0, alpha=[2.3, 0.64], label_smoothing=0.1)

class CustomModelTrainer(ClassificationTrainer):
"""
自定义训练器
重写模型创建方法，返回自定义的分类模型
"""
def get_model(self, cfg=None, weights=None, verbose=True):
"""
重写模型创建逻辑
:param cfg: 模型配置文件
:param weights: 预训练权重路径
:param verbose: 是否打印模型详情
:return: 自定义分类模型实例
"""
# 创建自定义模型实例
model = CustomClassificationModel(cfg, nc=self.data["nc"], ch=self.data["channels"], verbose=verbose)
# 加载预训练权重
if weights:
model.load(weights)
return model

def run_training():
# 训练超参数配置
overrides = dict(
model="yolo26s-cls.pt",
data="./cats_and_dogs_filtered",
imgsz=224,
epochs=5,
batch=8,
workers=2,
device=0,
amp=False,
patience=0,
save=False,
project="./runs/classify/cat_dog_standard",
name="custom_model",
exist_ok=True,
pretrained=True,
)

# 启动训练
trainer = CustomModelTrainer(overrides=overrides.copy())
trainer.train()

# 验证自定义损失生效情况
if hasattr(trainer.model, 'criterion') and hasattr(trainer.model.criterion, 'call_count'):
print(f"自定义损失函数累计调用次数: {trainer.model.criterion.call_count}")

if __name__ == "__main__":
run_training()