YOLOV8网络框架下.pt格式转ONNX的分析
·
目录
前言
视觉项目部署时，pt格式的模型文件推理速度较慢，而且占用内存较大，往往需要将其转换为ONNX格式，再根据实际的硬件情况，是否继续优化为OpenVINO或者TensorRT格式的模型文件。
本文主要分析一个pt文件如何转换为ONNX格式，并且使用转换的格式进行推理，会提供完整的代码，大家一起交流学习。
一、训练模型
笔者所用的训练文件，是基于Github上的博主开源的Pytorch版本的YOLOV8训练项目改进而来。
-->项目链接：bubbliiiing/yolov8-pytorch： 这是一个yolov8-pytorch的仓库，可以用于训练自己的数据集。
该博主还有其他版本的yolo训练项目，笔者在这里仅作说明和分享。
下面是我修改后的训练代码，调整了一些设置。
#-------------------------------------#
# 对数据集进行训练
#-------------------------------------#
import datetime
import os
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from functools import partial
import numpy as np
# 修改导入语句以使用新的模型
from model.YOLOV8 import YoloBody
from model.yolo_training import (Loss, ModelEMA, get_lr_scheduler,
set_optimizer_lr, weights_init)
from utils.callbacks import EvalCallback, LossHistory
from utils.dataloader import YoloDataset, yolo_dataset_collate
from utils.utils import (download_weights, get_classes, seed_everything,
show_config, worker_init_fn)
from utils.utils_fit import fit_one_epoch, print_detailed_metrics
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
if __name__ == "__main__":
#---------------------------------#
# Cuda 是否使用Cuda
# 没有GPU可以设置成False
#---------------------------------#
Cuda = True
#----------------------------------------------#
# Seed 用于固定随机种子
# 使得每次独立训练都可以获得一样的结果
#----------------------------------------------#
seed = 11
#---------------------------------------------------------------------#
# distributed 用于指定是否使用单机多卡分布式运行
# 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
# Windows系统下默认使用DP模式调用所有显卡，不支持DDP。
# DP模式：
# 设置 distributed = False
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py
# DDP模式：
# 设置 distributed = True
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
#---------------------------------------------------------------------#
distributed = False
#---------------------------------------------------------------------#
# sync_bn 是否使用sync_bn，DDP模式多卡可用
#---------------------------------------------------------------------#
sync_bn = False
#---------------------------------------------------------------------#
# fp16 是否使用混合精度训练
# 可减少约一半的显存、需要pytorch1.7.1以上
#---------------------------------------------------------------------#
fp16 = True
#---------------------------------------------------------------------#
# classes_path 指向model_data下的txt，与自己训练的数据集相关
# 训练前一定要修改classes_path，使其对应自己的数据集
#---------------------------------------------------------------------#
classes_path = 'model_data\\person_label.txt'
#----------------------------------------------------------------------------------------------------------------------------#
# 权值文件的下载请看README，可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的，因为特征是通用的。
# 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分，用于进行特征提取。
# 预训练权重对于99%的情况都必须要用，不用的话主干部分的权值太过随机，特征提取效果不明显，网络训练的结果也不会好
#
# 如果训练过程中存在中断训练的操作，可以将model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。
# 同时修改下方的 冻结阶段 或者 解冻阶段 的参数，来保证模型epoch的连续性。
#
# 当model_path = ''的时候不加载整个模型的权值。
#
# 此处使用的是整个模型的权重，因此是在train.py进行加载的。
# 如果想要让模型从0开始训练，则设置model_path = ''，下面的Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#
# 一般来讲，网络从0开始的训练效果会很差，因为权值太过随机，特征提取效果不明显，因此非常、非常、非常不建议大家从0开始训练！
# 从0开始训练有两个方案：
# 1、得益于Mosaic数据增强方法强大的数据增强能力，将UnFreeze_Epoch设置的较大（300及以上）、batch较大（16及以上）、数据较多（万以上）的情况下，
# 可以设置mosaic=True，直接随机初始化参数开始训练，但得到的效果仍然不如有预训练的情况。（像COCO这样的大数据集可以这样做）
# 2、了解imagenet数据集，首先训练分类模型，获得网络的主干部分权值，分类模型的 主干部分 和该模型通用，基于此进行训练。
#----------------------------------------------------------------------------------------------------------------------------#
model_path = r'model_data\\yolov8_s_backbone_weights.pth'
#------------------------------------------------------#
# input_shape 输入的shape大小，一定要是32的倍数
#------------------------------------------------------#
input_shape = [640, 640]
#------------------------------------------------------#
# phi 所使用到的yolov8的版本
# n : 对应yolov8_n
# s : 对应yolov8_s
# m : 对应yolov8_m
# l : 对应yolov8_l
# x : 对应yolov8_x
#------------------------------------------------------#
phi = 's'
#----------------------------------------------------------------------------------------------------------------------------#
# pretrained 是否使用主干网络的预训练权重，此处使用的是主干的权重，因此是在模型构建的时候进行加载的。
# 如果设置了model_path，则主干的权值无需加载，pretrained的值无意义。
# 如果不设置model_path，pretrained = True，此时仅加载主干开始训练。
# 如果不设置model_path，pretrained = False，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
pretrained = False
#------------------------------------------------------------------#
# mosaic 马赛克数据增强。
# mosaic_prob 每个step有多少概率使用mosaic数据增强，默认50%。
#
# mixup 是否使用mixup数据增强，仅在mosaic=True时有效。
# 只会对mosaic增强后的图片进行mixup的处理。
# mixup_prob 有多少概率在mosaic后使用mixup数据增强，默认50%。
# 总的mixup概率为mosaic_prob * mixup_prob。
#
# special_aug_ratio 参考YoloX，由于Mosaic生成的训练图片，远远脱离自然图片的真实分布。
# 当mosaic=True时，本代码会在special_aug_ratio范围内开启mosaic。
# 默认为前70%个epoch，100个世代会开启70个世代。
#------------------------------------------------------------------#
mosaic = True
mosaic_prob = 0.10
mixup = True
mixup_prob = 0.10
special_aug_ratio = 0.3
#------------------------------------------------------------------#
# label_smoothing 标签平滑。一般0.01以下。如0.01、0.005。
#------------------------------------------------------------------#
label_smoothing = 0.015
#----------------------------------------------------------------------------------------------------------------------------#
# 训练分为两个阶段，分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
# 冻结训练需要的显存较小，显卡非常差的情况下，可设置Freeze_Epoch等于UnFreeze_Epoch，Freeze_Train = True，此时仅仅进行冻结训练。
#
# 在此提供若干参数设置建议，各位训练者根据自己的需求进行灵活调整：
# （一）从整个模型的预训练权重开始训练：
# Adam：
# Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train = True，optimizer_type = 'adam'，Init_lr = 1e-3，weight_decay = 0。（冻结）
# Init_Epoch = 0，UnFreeze_Epoch = 100，Freeze_Train = False，optimizer_type = 'adam'，Init_lr = 1e-3，weight_decay = 0。（不冻结）
# SGD：
# Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 300，Freeze_Train = True，optimizer_type = 'sgd'，Init_lr = 1e-2，weight_decay = 5e-4。（冻结）
# Init_Epoch = 0，UnFreeze_Epoch = 300，Freeze_Train = False，optimizer_type = 'sgd'，Init_lr = 1e-2，weight_decay = 5e-4。（不冻结）
# 其中：UnFreeze_Epoch可以在100-300之间调整。
# （二）从0开始训练：
# Init_Epoch = 0，UnFreeze_Epoch >= 300，Unfreeze_batch_size >= 16，Freeze_Train = False（不冻结训练）
# 其中：UnFreeze_Epoch尽量不小于300。optimizer_type = 'sgd'，Init_lr = 1e-2，mosaic = True。
# （三）batch_size的设置：
# 在显卡能够接受的范围内，以大为好。显存不足与数据集大小无关，提示显存不足（OOM或者CUDA out of memory）请调小batch_size。
# 受到BatchNorm层影响，batch_size最小为2，不能为1。
# 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大，因为关系到学习率的自动调整。
#----------------------------------------------------------------------------------------------------------------------------#
#------------------------------------------------------------------#
# 冻结阶段训练参数
# 此时模型的主干被冻结了，特征提取网络不发生改变
# 占用的显存较小，仅对网络进行微调
# Init_Epoch 模型当前开始的训练世代，其值可以大于Freeze_Epoch，如设置：
# Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100
# 会跳过冻结阶段，直接从60代开始，并调整对应的学习率。
# （断点续练时使用）
# Freeze_Epoch 模型冻结训练的Freeze_Epoch
# (当Freeze_Train=False时失效)
# Freeze_batch_size 模型冻结训练的batch_size
# (当Freeze_Train=False时失效)
#------------------------------------------------------------------#
Init_Epoch = 0
Freeze_Epoch = 0
Freeze_batch_size = 24
#-----------------------------------------------------------------#
# 解冻阶段训练参数
# 此时模型的主干不被冻结了，特征提取网络会发生改变
# 占用的显存较大，网络所有的参数都会发生改变
# UnFreeze_Epoch 模型总共训练的epoch
# SGD需要更长的时间收敛，因此设置较大的UnFreeze_Epoch b
# Unfreeze_batch_size 模型在解冻后的batch_size
#------------------------------------------------------------------#
UnFreeze_Epoch = 400
Unfreeze_batch_size = 24
#------------------------------------------------------------------#
# Freeze_Train 是否进行冻结训练
# 默认先冻结主干训练后解冻训练。
#------------------------------------------------------------------#
Freeze_Train = True
#------------------------------------------------------------------#
# 其它训练参数：学习率、优化器、学习率下降有关
#------------------------------------------------------------------#
#------------------------------------------------------------------#
# Init_lr 模型的最大学习率
# Min_lr 模型的最小学习率，默认为最大学习率的0.01
#------------------------------------------------------------------#
Init_lr = 2e-2
Min_lr = Init_lr * 0.001
#------------------------------------------------------------------#
# optimizer_type 使用到的优化器种类，可选的有adam、sgd
# 当使用Adam优化器时建议设置 Init_lr=1e-3
# 当使用SGD优化器时建议设置 Init_lr=1e-2
# momentum 优化器内部使用到的momentum参数
# weight_decay 权值衰减，可防止过拟合
# adam会导致weight_decay错误，使用adam时建议设置为0。
#------------------------------------------------------------------#
optimizer_type = "sgd"
momentum = 0.9
weight_decay = 5e-4
#------------------------------------------------------------------#
# lr_decay_type 使用到的学习率下降方式，可选的有step、cos
#------------------------------------------------------------------#
lr_decay_type = "cos"
#------------------------------------------------------------------#
# save_period 多少个epoch保存一次权值
#------------------------------------------------------------------#
save_period = 10
#------------------------------------------------------------------#
# save_dir 权值与日志文件保存的文件夹
#------------------------------------------------------------------#
save_dir = 'logs'
#------------------------------------------------------------------#
# eval_flag 是否在训练时进行评估，评估对象为验证集
# 安装pycocotools库后，评估体验更佳。
# eval_period 代表多少个epoch评估一次，不建议频繁的评估
# 评估需要消耗较多的时间，频繁评估会导致训练非常慢
# 此处获得的mAP会与get_map.py获得的会有所不同，原因有二：
# （一）此处获得的mAP为验证集的mAP。
# （二）此处设置评估参数较为保守，目的是加快评估速度。
#------------------------------------------------------------------#
eval_flag = True
eval_period = 10
#------------------------------------------------------------------#
# num_workers 用于设置是否使用多线程读取数据
# 开启后会加快数据读取速度，但是会占用更多内存
# 内存较小的电脑可以设置为2或者0
#------------------------------------------------------------------#
num_workers = min(os.cpu_count(), 8)
#------------------------------------------------------#
# train_annotation_path 训练图片路径和标签
# val_annotation_path 验证图片路径和标签
#------------------------------------------------------#
train_annotation_path = '2025_train.txt'
val_annotation_path = '2025_val.txt'
#------------------------------------------------------------------#
# gradient_clip_norm 梯度裁剪的范数阈值，防止梯度爆炸
#------------------------------------------------------------------#
gradient_clip_norm = 5.0
#------------------------------------------------------------------#
# precision_save_params 精度导向保存机制参数
#------------------------------------------------------------------#
min_recall_threshold = 0.85
min_f1_threshold = 0.7
#------------------------------------------------------------------#
# export_deploy_model 是否额外导出一个用于 ONNX 转换的轻量模型
#------------------------------------------------------------------#
export_deploy_model = True # 推荐开启，方便后续 ONNX 导出
#==================================================#
# 开始执行训练流程
#==================================================#
seed_everything(seed)
#------------------------------------------------------#
# 设置用到的显卡
#------------------------------------------------------#
ngpus_per_node = torch.cuda.device_count()
if distributed:
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
rank = int(os.environ["RANK"])
device = torch.device("cuda", local_rank)
if local_rank == 0:
print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...")
print("Gpu Device Count : ", ngpus_per_node)
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
local_rank = 0
rank = 0
#------------------------------------------------------#
# 获取classes和anchor
#------------------------------------------------------#
class_names, num_classes = get_classes(classes_path)
#----------------------------------------------------#
# 下载预训练权重
#----------------------------------------------------#
if pretrained:
if distributed:
if local_rank == 0:
download_weights(phi)
dist.barrier()
else:
download_weights(phi)
#------------------------------------------------------#
# 创建yolo模型
#------------------------------------------------------#
model = YoloBody(input_shape, num_classes, phi, pretrained=pretrained)
if model_path != '':
if local_rank == 0:
print('Load weights {}.'.format(model_path))
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)
load_key, no_load_key, temp_dict = [], [], {}
for k, v in pretrained_dict.items():
if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
temp_dict[k] = v
load_key.append(k)
else:
no_load_key.append(k)
model_dict.update(temp_dict)
model.load_state_dict(model_dict)
if local_rank == 0:
print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))
print("\n\033[1;33;44m温馨提示，head部分没有载入是正常现象，Backbone部分没有载入是错误的。\033[0m")
#----------------------#
# 获得损失函数
#----------------------#
yolo_loss = Loss(model)
#----------------------#
# 记录Loss
#----------------------#
if local_rank == 0:
time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
log_dir = os.path.join(save_dir, "loss_" + str(time_str))
loss_history = LossHistory(log_dir, model, input_shape=input_shape)
else:
loss_history = None
#------------------------------------------------------------------#
# torch 1.2不支持amp，建议使用torch 1.7.1及以上正确使用fp16
#------------------------------------------------------------------#
if fp16:
from torch.cuda.amp import GradScaler
scaler = GradScaler()
else:
scaler = None
model_train = model.train()
#----------------------------#
# 多卡同步Bn
#----------------------------#
if sync_bn and ngpus_per_node > 1 and distributed:
model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train)
elif sync_bn:
print("Sync_bn is not support in one gpu or not distributed.")
if Cuda:
if distributed:
model_train = model_train.cuda(local_rank)
model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank], find_unused_parameters=True)
else:
model_train = torch.nn.DataParallel(model)
cudnn.benchmark = True
cudnn.deterministic = False
model_train = model_train.cuda()
#----------------------------#
# 权值平滑
#----------------------------#
ema = ModelEMA(model_train)
#---------------------------#
# 读取数据集对应的txt
#---------------------------#
with open(train_annotation_path, encoding='utf-8') as f:
train_lines = f.readlines()
with open(val_annotation_path, encoding='utf-8') as f:
val_lines = f.readlines()
num_train = len(train_lines)
num_val = len(val_lines)
if local_rank == 0:
show_config(
classes_path=classes_path, model_path=model_path, input_shape=input_shape,
Init_Epoch=Init_Epoch, Freeze_Epoch=Freeze_Epoch, UnFreeze_Epoch=UnFreeze_Epoch,
Freeze_batch_size=Freeze_batch_size, Unfreeze_batch_size=Unfreeze_batch_size,
Freeze_Train=Freeze_Train,
Init_lr=Init_lr, Min_lr=Min_lr, optimizer_type=optimizer_type, momentum=momentum,
lr_decay_type=lr_decay_type,
save_period=save_period, save_dir=save_dir, num_workers=num_workers,
num_train=num_train, num_val=num_val
)
wanted_step = 5e4 if optimizer_type == "sgd" else 1.5e4
total_step = num_train // Unfreeze_batch_size * UnFreeze_Epoch
if total_step <= wanted_step:
if num_train // Unfreeze_batch_size == 0:
raise ValueError('数据集过小，无法进行训练，请扩充数据集。')
wanted_epoch = wanted_step // (num_train // Unfreeze_batch_size) + 1
print("\n\033[1;33;44m[Warning] 使用%s优化器时，建议将训练总步长设置到%d以上。\033[0m" % (optimizer_type, wanted_step))
print("\033[1;33;44m[Warning] 本次运行的总训练数据量为%d，Unfreeze_batch_size为%d，共训练%d个Epoch，计算出总训练步长为%d。\033[0m" % (num_train, Unfreeze_batch_size, UnFreeze_Epoch, total_step))
print("\033[1;33;44m[Warning] 由于总训练步长为%d，小于建议总步长%d，建议设置总世代为%d。\033[0m" % (total_step, wanted_step, wanted_epoch))
#------------------------------------------------------#
# 开始训练
#------------------------------------------------------#
if True:
UnFreeze_flag = False
batch_size = Freeze_batch_size if Freeze_Train else Unfreeze_batch_size
nbs = 64
lr_limit_max = 1e-3 if optimizer_type == 'adam' else 5e-2
lr_limit_min = 3e-4 if optimizer_type == 'adam' else 5e-4
Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
pg0, pg1, pg2 = [], [], []
for k, v in model.named_modules():
if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias)
if isinstance(v, nn.BatchNorm2d) or "bn" in k:
pg0.append(v.weight)
elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight)
optimizer = {
'adam': optim.Adam(pg0, Init_lr_fit, betas=(momentum, 0.999)),
'sgd': optim.SGD(pg0, Init_lr_fit, momentum=momentum, nesterov=True)
}[optimizer_type]
optimizer.add_param_group({"params": pg1, "weight_decay": weight_decay})
optimizer.add_param_group({"params": pg2})
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
if ema:
ema.updates = epoch_step * Init_Epoch
train_dataset = YoloDataset(train_lines, input_shape, num_classes, epoch_length=UnFreeze_Epoch,
mosaic=mosaic, mixup=mixup, mosaic_prob=mosaic_prob, mixup_prob=mixup_prob,
train=True, special_aug_ratio=special_aug_ratio)
val_dataset = YoloDataset(val_lines, input_shape, num_classes, epoch_length=UnFreeze_Epoch,
mosaic=False, mixup=False, mosaic_prob=0, mixup_prob=0,
train=False, special_aug_ratio=0)
if distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True)
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False)
batch_size = batch_size // ngpus_per_node
shuffle = False
else:
train_sampler = None
val_sampler = None
shuffle = True
gen = DataLoader(train_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers,
pin_memory=True, drop_last=True, collate_fn=yolo_dataset_collate,
sampler=train_sampler,
worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed),
persistent_workers=True if num_workers > 0 else False,
prefetch_factor=3 if num_workers > 0 else None)
gen_val = DataLoader(val_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers,
pin_memory=True, drop_last=True, collate_fn=yolo_dataset_collate,
sampler=val_sampler,
worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed),
persistent_workers=True if num_workers > 0 else False,
prefetch_factor=3 if num_workers > 0 else None)
if local_rank == 0:
eval_callback = EvalCallback(model, input_shape, class_names, num_classes, val_lines, log_dir, Cuda,
eval_flag=eval_flag, period=eval_period, verbose=False)
else:
eval_callback = None
# 初始化最佳指标
best_map = 0.0
best_precision = 0.0
best_recall = 0.0
best_f1 = 0.0
if model_path != '' and os.path.exists(model_path):
try:
checkpoint = torch.load(model_path, map_location=device)
if 'best_map' in checkpoint:
best_map = checkpoint['best_map']
print(f"Loaded previous best mAP: {best_map:.4f}")
if 'best_precision' in checkpoint:
best_precision = checkpoint['best_precision']
print(f"Loaded previous best precision: {best_precision:.4f}")
if 'best_recall' in checkpoint:
best_recall = checkpoint['best_recall']
print(f"Loaded previous best recall: {best_recall:.4f}")
if 'best_f1' in checkpoint:
best_f1 = checkpoint['best_f1']
print(f"Loaded previous best F1-score: {best_f1:.4f}")
except Exception as e:
print(f"Could not load best metrics: {e}")
for epoch in range(Init_Epoch, UnFreeze_Epoch):
if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train:
batch_size = Unfreeze_batch_size
Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
for param in model.backbone.parameters():
param.requires_grad = True
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
if ema:
ema.updates = epoch_step * epoch
if distributed:
batch_size = batch_size // ngpus_per_node
gen = DataLoader(train_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers,
pin_memory=True, drop_last=True, collate_fn=yolo_dataset_collate,
sampler=train_sampler,
worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed),
persistent_workers=True if num_workers > 0 else False,
prefetch_factor=2 if num_workers > 0 else None)
gen_val = DataLoader(val_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers,
pin_memory=True, drop_last=True, collate_fn=yolo_dataset_collate,
sampler=val_sampler,
worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed),
persistent_workers=True if num_workers > 0 else False,
prefetch_factor=2 if num_workers > 0 else None)
UnFreeze_flag = True
gen.dataset.epoch_now = epoch
gen_val.dataset.epoch_now = epoch
if distributed:
train_sampler.set_epoch(epoch)
if epoch > UnFreeze_Epoch * 0.7:
decay_factor = max(0.1, 1.0 - (epoch - UnFreeze_Epoch * 0.7) / (UnFreeze_Epoch * 0.3))
current_mosaic_prob = mosaic_prob * decay_factor
current_mixup_prob = mixup_prob * decay_factor
gen.dataset.mosaic_prob = current_mosaic_prob
gen.dataset.mixup_prob = current_mixup_prob
if local_rank == 0 and epoch % 10 == 0:
print(f"Epoch {epoch}: Mosaic={current_mosaic_prob:.3f}, Mixup={current_mixup_prob:.3f}")
set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
fit_result = fit_one_epoch(model_train, model, ema, yolo_loss, loss_history, eval_callback,
optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val,
UnFreeze_Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rank)
if local_rank == 0 and fit_result is not None:
print(f'Epoch:{epoch + 1}/{UnFreeze_Epoch} | Train time: {fit_result["train_time"]:.2f}s '
f'| Val time: {fit_result["val_time"]:.2f}s | Train loss: {fit_result["train_loss"]:.4f} '
f'| Val loss: {fit_result["val_loss"]:.4f} | LR: {get_lr(optimizer):.6f}')
# 评估并判断是否更新最佳模型
should_save_checkpoint = (epoch + 1) % save_period == 0 or (epoch + 1) == UnFreeze_Epoch
current_map = 0
current_precision = 0
current_recall = 0
current_f1 = 0
if (epoch + 1) % eval_period == 0 and hasattr(eval_callback, 'metrics') and eval_callback.metrics:
print_detailed_metrics(eval_callback.metrics)
current_map = eval_callback.metrics.get('mAP', 0)
current_precision = eval_callback.metrics.get('precision', 0)
current_recall = eval_callback.metrics.get('recall', 0)
if current_precision + current_recall > 0:
current_f1 = 2 * (current_precision * current_recall) / (current_precision + current_recall)
else:
current_f1 = 0
print(f"🔍 Epoch {epoch + 1}: Precision={current_precision:.4f}, Recall={current_recall:.4f}, F1-Score={current_f1:.4f}")
if current_recall >= min_recall_threshold and current_precision > best_precision:
print(f'\033[1;32m🏆 Epoch {epoch + 1}: Precision improved from {best_precision:.4f} to {current_precision:.4f}!\033[0m')
best_precision = current_precision
if current_precision >= 0.8 and current_recall > best_recall:
print(f'\033[1;32m🔍 Epoch {epoch + 1}: Recall improved from {best_recall:.4f} to {current_recall:.4f}!\033[0m')
best_recall = current_recall
if current_recall >= min_recall_threshold and current_f1 > best_f1 and current_f1 >= min_f1_threshold:
print(f'\033[1;32m🎯 Epoch {epoch + 1}: F1-score improved from {best_f1:.4f} to {current_f1:.4f}!\033[0m')
best_f1 = current_f1
if current_map > best_map:
print(f'\033[1;32m🏆 Epoch {epoch + 1}: mAP improved from {best_map:.4f} to {current_map:.4f}!\033[0m')
best_map = current_map
# ========== 保存完整 Checkpoint（用于断点续训）==========
if should_save_checkpoint:
ckpt_path = os.path.join(save_dir, f"ep{epoch + 1}_ckpt.pt")
save_dict = {
'model': model.module.state_dict() if hasattr(model, 'module') else model.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch + 1,
'ema': ema.ema.state_dict() if ema else None,
'best_map': best_map,
'best_precision': best_precision,
'best_recall': best_recall,
'best_f1': best_f1,
# --- 元信息 ---
'input_shape': input_shape,
'num_classes': num_classes,
'phi': phi,
'classes_path': classes_path
}
torch.save(save_dict, ckpt_path)
print(f'\033[1;34m💾 Full checkpoint saved to {ckpt_path}\033[0m')
# ========== 额外导出轻量部署模型（仅推理用）==========
if export_deploy_model and current_precision == best_precision:
deploy_path = os.path.join(save_dir, "model_best_precision_deploy.pt")
deploy_dict = {
'model': ema.ema.state_dict() if ema else save_dict['model'], # 使用 EMA 模型
'input_shape': input_shape, # 输入图片的尺寸
'num_classes': num_classes, # 类别数量
'phi': phi # 模型参数
}
torch.save(deploy_dict, deploy_path)
print(f"\033[1;32m🚀 Deploy model updated at: {deploy_path}\033[0m")
if Cuda and epoch % 5 == 0:
torch.cuda.empty_cache()
if distributed:
dist.barrier()
if local_rank == 0:
loss_history.writer.close()
在此处贴出代码的用意是，我们在训练模型时，会设定保存的模型参数，一些参数会影响到ONNX的格式转换。
我的代码中，保存的模型格式如下：
这里保存了两个模型，一个是用于断点续训练，一个是做推理使用，而我们转ONNX时需要使用推理使用的文件。
为什么要设置两种保存模型呢？
因为完整的pt文件，为了满足断点续训的要求，会保存完整的训练信息，会将模型权重，优化器状态，训练轮次等等信息保存，而这些信息是不在ONNX使用的，移除这些不必要的元数据，可以提高转换效率，也可以避免训练相关状态带来的转换问题。
.pt和.pth的区别？
这两种格式在本质上没有区别，但是在使用习惯和约定上有一些差异。
.pt文件，通常保存完整的模型检查点，包含模型权重、优化器状态、训练轮次等完整训练信息
.pth文件，传统上更多用于保存预训练模型权重
不过我感觉直接用pth文件做续训效果好一点
二、转ONNX
我先贴出我的转换脚本。
# export_onnx.py
import torch
import os
# 注意：确保 model.YOLOV8 与你的项目结构一致
from model.YOLOV8 import YoloBody
class YoloWrapper(torch.nn.Module):
"""包装器模型，只输出box和cls，便于ONNX导出"""
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
outputs = self.model(x) # 获取原始输出
# 只返回box和cls
return outputs[0], outputs[1]
def print_output_shapes(outputs, prefix="输出"):
"""递归打印输出形状的辅助函数"""
if isinstance(outputs, torch.Tensor):
print(f" {prefix} 形状: {outputs.shape}")
elif isinstance(outputs, (list, tuple)):
for i, item in enumerate(outputs):
print_output_shapes(item, f"{prefix}[{i}]")
else:
print(f" {prefix} 类型: {type(outputs)}")
def main():
# -------------------------------
# 模型路径（由训练脚本生成）
# -------------------------------
deploy_model_path = "logs/model_best_precision_deploy.pt"
if not os.path.exists(deploy_model_path):
raise FileNotFoundError(f"未找到部署模型文件: {deploy_model_path}")
print(f"Loading deployment model from {deploy_model_path}")
ckpt = torch.load(deploy_model_path, map_location='cpu')
# 自动提取参数
input_shape = ckpt['input_shape'] # 如 [640, 640]
num_classes = ckpt['num_classes']
phi = ckpt['phi']
print(f"Model Config: input_shape={input_shape}, num_classes={num_classes}, phi={phi}")
# 构建模型
model = YoloBody(
input_shape=input_shape,
num_classes=num_classes,
phi=phi,
pretrained=False
)
model.load_state_dict(ckpt['model']) #加载权重
model.eval() #设置为评估模式
# 测试原始模型输出
print("\n测试原始模型输出...")
dummy_input = torch.randn(1, 3, input_shape[0], input_shape[1])
with torch.no_grad():
original_outputs = model(dummy_input)
print(f"原始模型输出数量: {len(original_outputs) if isinstance(original_outputs, (list, tuple)) else 1}")
print_output_shapes(original_outputs, "原始输出")
# 创建包装器模型
wrapper_model = YoloWrapper(model)
wrapper_model.eval()
# 创建 dummy input
dummy_input = torch.randn(1, 3, input_shape[0], input_shape[1])
# 测试包装器模型输出
print("\n测试包装器模型输出...")
with torch.no_grad():
wrapper_outputs = wrapper_model(dummy_input)
print(f"包装器模型输出数量: {len(wrapper_outputs) if isinstance(wrapper_outputs, (list, tuple)) else 1}")
print_output_shapes(wrapper_outputs, "包装器输出")
# 输出路径
onnx_output = "yolov8_person_simple.onnx"
# 导出 ONNX
print(f"\n开始导出简化版ONNX模型...")
torch.onnx.export(
wrapper_model,
dummy_input,
onnx_output,
input_names=["images"],
output_names=["box", "cls"], # 只输出两个
dynamic_axes={
"images": {0: "batch"},
"box": {0: "batch", 2: "anchors"},
"cls": {0: "batch", 2: "anchors"}
},
opset_version=13,
do_constant_folding=True,
verbose=True, # 开启详细输出
export_params=True # 确保权重嵌入
)
print(f"✅ 成功导出简化版 ONNX 模型: {onnx_output}")
# ==========================
# 可选：使用 onnx-simplifier 进一步优化
# ==========================
try:
import onnx
from onnxsim import simplify
print("正在简化 ONNX 模型...")
onnx_model = onnx.load(onnx_output)
simplified_model, check = simplify(onnx_model)
if check:
simplified_output = "yolov8_person_simple_sim.onnx"
onnx.save(simplified_model, simplified_output)
print(f"✅ 简化完成 → {simplified_output}")
else:
print("⚠️ 简化失败，保留原始 ONNX")
except ImportError:
print("💡 提示：安装 onnxsim 可进一步优化模型: pip install onnxsim")
print("\n🔧 重要提示:")
print("1. 简化版ONNX模型只包含box和cls输出")
print("2. 在predict.py中需要将model_path改为新的ONNX模型")
print("3. 解码函数会自动处理简化版输出")
if __name__ == "__main__":
main()这个是网络框架文件：
import numpy as np
import torch
import torch.nn as nn
from model.yolo_training import weights_init
from utils.utils_bbox import make_anchors
def autopad(k, p=None, d=1):
# kernel, padding, dilation
# 对输入的特征层进行自动padding，按照Same原则
if d > 1:
# actual kernel-size
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]
if p is None:
# auto-pad
p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
return p
class SiLU(nn.Module):
# SiLU激活函数
@staticmethod
def forward(x):
return x * torch.sigmoid(x)
class Conv(nn.Module):
# 标准卷积+标准化+激活函数
default_act = SiLU()
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self, x):
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
return self.act(self.conv(x))
class Bottleneck(nn.Module):
# 标准瓶颈结构，残差结构
# c1为输入通道数，c2为输出通道数
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, k[0], 1)
self.cv2 = Conv(c_, c2, k[1], 1, g=g)
self.add = shortcut and c1 == c2
def forward(self, x):
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
class C2f(nn.Module):
# CSPNet结构结构，大残差结构
# c1为输入通道数，c2为输出通道数
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
super().__init__()
self.c = int(c2 * e)
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1)
self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
def forward(self, x):
# 进行一个卷积，然后划分成两份，每个通道都为c
y = list(self.cv1(x).split((self.c, self.c), 1))
# 每进行一次残差结构都保留，然后堆叠在一起，密集残差
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
class SPPF(nn.Module):
# SPP结构，5、9、13最大池化核的最大池化。
def __init__(self, c1, c2, k=5):
super().__init__()
c_ = c1 // 2
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
class Backbone(nn.Module):
def __init__(self, base_channels, base_depth, deep_mul, phi, pretrained=False):
super().__init__()
#-----------------------------------------------#
# 输入图片是3, 640, 640
#-----------------------------------------------#
# 3, 640, 640 => 32, 640, 640 => 64, 320, 320
self.stem = Conv(3, base_channels, 3, 2)
# 64, 320, 320 => 128, 160, 160 => 128, 160, 160
self.dark2 = nn.Sequential(
Conv(base_channels, base_channels * 2, 3, 2),
C2f(base_channels * 2, base_channels * 2, base_depth, True),
)
# 128, 160, 160 => 256, 80, 80 => 256, 80, 80
self.dark3 = nn.Sequential(
Conv(base_channels * 2, base_channels * 4, 3, 2),
C2f(base_channels * 4, base_channels * 4, base_depth * 2, True),
)
# 256, 80, 80 => 512, 40, 40 => 512, 40, 40
self.dark4 = nn.Sequential(
Conv(base_channels * 4, base_channels * 8, 3, 2),
C2f(base_channels * 8, base_channels * 8, base_depth * 2, True),
)
# 512, 40, 40 => 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 20, 20
self.dark5 = nn.Sequential(
Conv(base_channels * 8, int(base_channels * 16 * deep_mul), 3, 2),
C2f(int(base_channels * 16 * deep_mul), int(base_channels * 16 * deep_mul), base_depth, True),
SPPF(int(base_channels * 16 * deep_mul), int(base_channels * 16 * deep_mul), k=5)
)
if pretrained:
url = {
"n" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_n_backbone_weights.pth',
"s" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_s_backbone_weights.pth',
"m" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_m_backbone_weights.pth',
"l" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_l_backbone_weights.pth',
"x" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_x_backbone_weights.pth',
}[phi]
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", model_dir="./model_data")
self.load_state_dict(checkpoint, strict=False)
print("Load weights from " + url.split('/')[-1])
def forward(self, x):
x = self.stem(x)
x = self.dark2(x)
#-----------------------------------------------#
# dark3的输出为256, 80, 80，是一个有效特征层
#-----------------------------------------------#
x = self.dark3(x)
feat1 = x
#-----------------------------------------------#
# dark4的输出为512, 40, 40，是一个有效特征层
#-----------------------------------------------#
x = self.dark4(x)
feat2 = x
#-----------------------------------------------#
# dark5的输出为1024 * deep_mul, 20, 20，是一个有效特征层
#-----------------------------------------------#
x = self.dark5(x)
feat3 = x
return feat1, feat2, feat3
def fuse_conv_and_bn(conv, bn):
# 混合Conv2d + BatchNorm2d 减少计算量
# Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
fusedconv = nn.Conv2d(conv.in_channels,
conv.out_channels,
kernel_size=conv.kernel_size,
stride=conv.stride,
padding=conv.padding,
dilation=conv.dilation,
groups=conv.groups,
bias=True).requires_grad_(False).to(conv.weight.device)
# 准备kernel
w_conv = conv.weight.clone().view(conv.out_channels, -1)
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
# 准备bias
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
return fusedconv
class DFL(nn.Module):
# DFL模块
# Distribution Focal Loss (DFL) proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
def __init__(self, c1=16):
super().__init__()
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(c1, dtype=torch.float)
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
self.c1 = c1
def forward(self, x):
# bs, self.reg_max * 4, 8400
b, c, a = x.shape
# bs, 4, self.reg_max, 8400 => bs, self.reg_max, 4, 8400 => b, 4, 8400
# 以softmax的方式，对0~16的数字计算百分比，获得最终数字。
return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
# return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
#---------------------------------------------------#
# yolo_body
#---------------------------------------------------#
class YoloBody(nn.Module):
def __init__(self, input_shape, num_classes, phi, pretrained=False):
super(YoloBody, self).__init__()
depth_dict = {'n' : 0.33, 's' : 0.33, 'm' : 0.67, 'l' : 1.00, 'x' : 1.00,}
width_dict = {'n' : 0.25, 's' : 0.50, 'm' : 0.75, 'l' : 1.00, 'x' : 1.25,}
deep_width_dict = {'n' : 1.00, 's' : 1.00, 'm' : 0.75, 'l' : 0.50, 'x' : 0.50,}
dep_mul, wid_mul, deep_mul = depth_dict[phi], width_dict[phi], deep_width_dict[phi]
base_channels = int(wid_mul * 64) # 64
base_depth = max(round(dep_mul * 3), 1) # 3
#-----------------------------------------------#
# 输入图片是3, 640, 640
#-----------------------------------------------#
#---------------------------------------------------#
# 生成主干模型
# 获得三个有效特征层，他们的shape分别是：
# 256, 80, 80
# 512, 40, 40
# 1024 * deep_mul, 20, 20
#---------------------------------------------------#
self.backbone = Backbone(base_channels, base_depth, deep_mul, phi, pretrained=pretrained)
#------------------------加强特征提取网络------------------------#
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
self.conv3_for_upsample1 = C2f(int(base_channels * 16 * deep_mul) + base_channels * 8, base_channels * 8, base_depth, shortcut=False)
# 768, 80, 80 => 256, 80, 80
self.conv3_for_upsample2 = C2f(base_channels * 8 + base_channels * 4, base_channels * 4, base_depth, shortcut=False)
# 256, 80, 80 => 256, 40, 40
self.down_sample1 = Conv(base_channels * 4, base_channels * 4, 3, 2)
# 512 + 256, 40, 40 => 512, 40, 40
self.conv3_for_downsample1 = C2f(base_channels * 8 + base_channels * 4, base_channels * 8, base_depth, shortcut=False)
# 512, 40, 40 => 512, 20, 20
self.down_sample2 = Conv(base_channels * 8, base_channels * 8, 3, 2)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
self.conv3_for_downsample2 = C2f(int(base_channels * 16 * deep_mul) + base_channels * 8, int(base_channels * 16 * deep_mul), base_depth, shortcut=False)
#------------------------加强特征提取网络------------------------#
ch = [base_channels * 4, base_channels * 8, int(base_channels * 16 * deep_mul)]
self.shape = None
self.nl = len(ch)
# self.stride = torch.zeros(self.nl)
self.stride = torch.tensor([256 / x.shape[-2] for x in self.backbone.forward(torch.zeros(1, 3, 256, 256))]) # forward
self.reg_max = 16 # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
self.no = num_classes + self.reg_max * 4 # number of outputs per anchor
self.num_classes = num_classes
c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], num_classes) # channels
self.cv2 = nn.ModuleList(nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, num_classes, 1)) for x in ch)
if not pretrained:
weights_init(self)
self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
def fuse(self):
print('Fusing layers... ')
for m in self.modules():
if type(m) is Conv and hasattr(m, 'bn'):
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
delattr(m, 'bn') # remove batchnorm
m.forward = m.forward_fuse # update forward
return self
def forward(self, x):
# backbone
feat1, feat2, feat3 = self.backbone.forward(x)
#------------------------加强特征提取网络------------------------#
# 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 40, 40
P5_upsample = self.upsample(feat3)
# 1024 * deep_mul, 40, 40 cat 512, 40, 40 => 1024 * deep_mul + 512, 40, 40
P4 = torch.cat([P5_upsample, feat2], 1)
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
P4 = self.conv3_for_upsample1(P4)
# 512, 40, 40 => 512, 80, 80
P4_upsample = self.upsample(P4)
# 512, 80, 80 cat 256, 80, 80 => 768, 80, 80
P3 = torch.cat([P4_upsample, feat1], 1)
# 768, 80, 80 => 256, 80, 80
P3 = self.conv3_for_upsample2(P3)
# 256, 80, 80 => 256, 40, 40
P3_downsample = self.down_sample1(P3)
# 512, 40, 40 cat 256, 40, 40 => 768, 40, 40
P4 = torch.cat([P3_downsample, P4], 1)
# 768, 40, 40 => 512, 40, 40
P4 = self.conv3_for_downsample1(P4)
# 512, 40, 40 => 512, 20, 20
P4_downsample = self.down_sample2(P4)
# 512, 20, 20 cat 1024 * deep_mul, 20, 20 => 1024 * deep_mul + 512, 20, 20
P5 = torch.cat([P4_downsample, feat3], 1)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
P5 = self.conv3_for_downsample2(P5)
#------------------------加强特征提取网络------------------------#
# P3 256, 80, 80
# P4 512, 40, 40
# P5 1024 * deep_mul, 20, 20
shape = P3.shape # BCHW
# P3 256, 80, 80 => num_classes + self.reg_max * 4, 80, 80
# P4 512, 40, 40 => num_classes + self.reg_max * 4, 40, 40
# P5 1024 * deep_mul, 20, 20 => num_classes + self.reg_max * 4, 20, 20
x = [P3, P4, P5]
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
# num_classes + self.reg_max * 4 , 8400 => cls num_classes, 8400;
# box self.reg_max * 4, 8400
box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.num_classes), 1)
# origin_cls = [xi.split((self.reg_max * 4, self.num_classes), 1)[1] for xi in x]
dbox = self.dfl(box)
return dbox, cls, x, self.anchors.to(dbox.device), self.strides.to(dbox.device)以上两份代码，便于我们更好的了解转换的流程，接下来我们对模型的转换做一个简单的分析。
2.1、提取pt格式下模型的参数
这一步提取的参数，就是训练代码保存的那几个参数。
2.2、构建模型
这一步的作用是调用YOLO的网络框架，方便后续参数的接收和保存
2.3、基于构建的模型包装新模型
包装的模型只接收YOLO网络框架中的box和cls，即检测框信息和类别置信度信息。
如何获取到的呢？
我们来看YOLO网络框架内的代码，在前向传播的过程中，会对输入的参数进行一系列的计算，最后会返回检测到的数据，而我们仅需要接收box和cls。
def forward(self, x):
# backbone
feat1, feat2, feat3 = self.backbone.forward(x)
#------------------------加强特征提取网络------------------------#
# 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 40, 40
P5_upsample = self.upsample(feat3)
# 1024 * deep_mul, 40, 40 cat 512, 40, 40 => 1024 * deep_mul + 512, 40, 40
P4 = torch.cat([P5_upsample, feat2], 1)
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
P4 = self.conv3_for_upsample1(P4)
# 512, 40, 40 => 512, 80, 80
P4_upsample = self.upsample(P4)
# 512, 80, 80 cat 256, 80, 80 => 768, 80, 80
P3 = torch.cat([P4_upsample, feat1], 1)
# 768, 80, 80 => 256, 80, 80
P3 = self.conv3_for_upsample2(P3)
# 256, 80, 80 => 256, 40, 40
P3_downsample = self.down_sample1(P3)
# 512, 40, 40 cat 256, 40, 40 => 768, 40, 40
P4 = torch.cat([P3_downsample, P4], 1)
# 768, 40, 40 => 512, 40, 40
P4 = self.conv3_for_downsample1(P4)
# 512, 40, 40 => 512, 20, 20
P4_downsample = self.down_sample2(P4)
# 512, 20, 20 cat 1024 * deep_mul, 20, 20 => 1024 * deep_mul + 512, 20, 20
P5 = torch.cat([P4_downsample, feat3], 1)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
P5 = self.conv3_for_downsample2(P5)
#------------------------加强特征提取网络------------------------#
# P3 256, 80, 80
# P4 512, 40, 40
# P5 1024 * deep_mul, 20, 20
shape = P3.shape # BCHW
# P3 256, 80, 80 => num_classes + self.reg_max * 4, 80, 80
# P4 512, 40, 40 => num_classes + self.reg_max * 4, 40, 40
# P5 1024 * deep_mul, 20, 20 => num_classes + self.reg_max * 4, 20, 20
x = [P3, P4, P5]
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
# num_classes + self.reg_max * 4 , 8400 => cls num_classes, 8400;
# box self.reg_max * 4, 8400
box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.num_classes), 1)
# origin_cls = [xi.split((self.reg_max * 4, self.num_classes), 1)[1] for xi in x]
dbox = self.dfl(box)
return dbox, cls, x, self.anchors.to(dbox.device), self.strides.to(dbox.device)2.4、转换格式
wrapper_model: 要导出的模型（这里是包装器模型，只输出box和cls）
dummy_input: 虚拟输入张量，用于推断模型结构和张量形状
onnx_output: 输出ONNX文件路径（yolov8_person_simple.onnx）
input_names=["images"]: 指定输入节点名称为"images"
output_names=["box", "cls"]: 指定输出节点名称，只保留两个主要输出
opset_version=13: 使用ONNX opset版本13，兼容性较好
do_constant_folding=True: 执行常量折叠优化，减小模型体积
verbose=True: 显示详细导出过程信息
export_params=True: 导出模型参数（权重），生成可直接使用的模型
至此即将pt格式转换为ONNX格式。
建议下载的onnxruntime要支持gpu加速的版本，检测的速度可以比pytorch格式的快两到三倍。
三、ONNX推理
推理文件：
# ONNX_predict.py
import torch
import numpy as np
from PIL import Image
import onnxruntime as ort
import os
# 添加ONNX Runtime支持检查
try:
import onnxruntime as ort
ONNX_AVAILABLE = True
except ImportError:
print("ONNX Runtime未安装，请使用: pip install onnxruntime")
ONNX_AVAILABLE = False
class ONNX_Detector:
"""
ONNX模型检测器类，用于处理ONNX格式的目标检测模型
"""
def __init__(self, model_path, num_classes, input_shape=(640, 640)):
"""
初始化ONNX检测器
Args:
model_path (str): ONNX模型文件路径
num_classes (int): 检测类别数
input_shape (tuple): 输入图像尺寸(height, width)
"""
if not ONNX_AVAILABLE:
raise ImportError("ONNX Runtime不可用，请安装: pip install onnxruntime")
if not os.path.exists(model_path):
raise FileNotFoundError(f"ONNX模型文件不存在: {model_path}")
# 初始化ONNX模型
self.model_path = model_path
self.num_classes = num_classes
self.input_shape = input_shape
# 配置ONNX Runtime以支持CUDA
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
try:
self.session = ort.InferenceSession(model_path, providers=providers)
print(f"✅ ONNX模型加载成功: {model_path} (GPU加速)")
except:
self.session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
print(f"✅ ONNX模型加载成功: {model_path} (CPU模式)")
self.input_name = self.session.get_inputs()[0].name
self.output_info = [(out.name, out.shape) for out in self.session.get_outputs()]
print(f"ONNX模型输出信息: {self.output_info}")
# 初始化解码器
from utils.utils_bbox import DecodeBox
self.decodebox = DecodeBox(num_classes=num_classes, input_shape=input_shape)
def preprocess_image(self, image):
"""
图像预处理
Args:
image (PIL.Image): 输入图像
Returns:
tuple: (处理后的图像tensor, 原始图像尺寸, 原始图像)
"""
image_shape = np.array(np.shape(image)[0:2])
# 转换为RGB格式
image = image.convert('RGB')
# 保存原始尺寸
original_width, original_height = image.size
# 调整图像大小（保持宽高比的letterbox）
scale = min(self.input_shape[0] / original_height, self.input_shape[1] / original_width)
new_width = int(original_width * scale)
new_height = int(original_height * scale)
# 创建新的图像并粘贴调整后的图像
resized_img = image.resize((new_width, new_height), Image.BICUBIC)
new_image = Image.new('RGB', self.input_shape, (128, 128, 128))
new_image.paste(resized_img, ((self.input_shape[1] - new_width) // 2, (self.input_shape[0] - new_height) // 2))
# 转换为numpy数组并归一化
image_data = np.array(new_image, dtype='float32') / 255.0
# 转换为CHW格式
image_data = np.transpose(image_data, (2, 0, 1))
# 添加批次维度
image_data = np.expand_dims(image_data, axis=0)
# 转换为tensor
images = torch.from_numpy(image_data)
return images, image_shape, image
def detect(self, image, conf_thres=0.5, nms_thres=0.3):
"""
对图像进行目标检测
Args:
image (PIL.Image): 输入图像
conf_thres (float): 置信度阈值
nms_thres (float): NMS阈值
Returns:
list: 检测结果列表，每个元素为[x0, y0, x1, y1, score, class_id]
"""
# 图像预处理
images, image_shape, original_image = self.preprocess_image(image)
# ONNX模型推理
try:
outputs = self.session.run(None, {self.input_name: images.cpu().numpy()})
outputs = [torch.from_numpy(output) for output in outputs]
except Exception as e:
print(f"ONNX模型推理错误: {e}")
return []
# 解码检测结果
try:
results = self.decodebox.decode_onnx_box(outputs)
except Exception as e:
print(f"解码错误: {e}")
return []
# 应用非极大值抑制
try:
results = self.decodebox.non_max_suppression(
results,
self.num_classes,
input_shape=self.input_shape,
image_shape=image_shape,
letterbox_image=True,
conf_thres=conf_thres,
nms_thres=nms_thres
)
except Exception as e:
print(f"NMS处理错误: {e}")
return []
# 处理检测结果
detections = []
if results[0] is not None:
top_label = np.array(results[0][:, 5], dtype='int32')
top_conf = results[0][:, 4]
top_boxes = results[0][:, :4]
for i in range(len(top_label)):
if top_conf[i] < conf_thres:
continue
box = top_boxes[i]
top, left, bottom, right = box
# 确保坐标在有效范围内
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(np.shape(original_image)[0], np.floor(bottom).astype('int32'))
right = min(np.shape(original_image)[1], np.floor(right).astype('int32'))
# 确保坐标顺序正确
x0 = min(left, right)
x1 = max(left, right)
y0 = min(top, bottom)
y1 = max(top, bottom)
if x1 > x0 and y1 > y0:
detections.append([x0, y0, x1, y1, top_conf[i], top_label[i]])
return detections
def draw_detections(self, image, detections, class_names=None):
"""
在图像上绘制检测结果
Args:
image (PIL.Image): 原始图像
detections (list): 检测结果列表
class_names (list): 类别名称列表
Returns:
PIL.Image: 绘制了检测结果的图像
"""
from PIL import ImageDraw, ImageFont
import numpy as np
# 复制原图用于绘制结果
result_image = image.copy()
draw = ImageDraw.Draw(result_image)
# 设置字体
try:
font_size = max(16, int(np.floor(3e-2 * np.shape(image)[1] + 10)))
font = ImageFont.truetype(font='model_data/simhei.ttf', size=font_size)
except:
try:
font = ImageFont.truetype(font='arial.ttf', size=font_size)
except:
font = ImageFont.load_default()
thickness = max(2, int(max((np.shape(image)[0] + np.shape(image)[1]) // self.input_shape[0], 1)))
# 绘制检测框
for i, det in enumerate(detections):
if len(det) < 6:
continue
x0, y0, x1, y1, score, class_id = det[:6]
# 使用类别名称
predicted_class = class_names[int(class_id)] if class_names and int(class_id) < len(class_names) else f"Class {int(class_id)}"
# 绘制边界框
color = (255, 0, 0)
x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
if x1 > x0 and y1 > y0:
draw.rectangle([x0, y0, x1, y1], outline=color, width=thickness)
label = f'{predicted_class} {score:.2f}'
draw.text((x0, y0), label, fill=color, font=font)
return result_image
# 使用示例
if __name__ == "__main__":
# 示例使用方法
try:
detector = ONNX_Detector(
model_path="yolov8_person_simple_sim.onnx",
num_classes=1,
input_shape=(640, 640)
)
# 加载测试图像
image = Image.open("test_image.jpg")
# 进行检测
detections = detector.detect(image, conf_thres=0.5, nms_thres=0.3)
print(f"检测到 {len(detections)} 个目标:")
for i, det in enumerate(detections):
x0, y0, x1, y1, score, class_id = det
print(f"目标 {i}: 边界框=({x0}, {y0}, {x1}, {y1}), 置信度={score:.2f}, 类别={class_id}")
# 绘制检测结果
class_names = ['person'] # 根据实际情况修改
result_image = detector.draw_detections(image, detections, class_names)
result_image.show()
except Exception as e:
print(f"发生错误: {e}")检测框解码文件：
import numpy as np
import torch
from torchvision.ops import nms
import pkg_resources as pkg
def check_version(current: str = "0.0.0",
minimum: str = "0.0.0",
name: str = "version ",
pinned: bool = False) -> bool:
current, minimum = (pkg.parse_version(x) for x in (current, minimum))
result = (current == minimum) if pinned else (current >= minimum) # bool
return result
TORCH_1_10 = check_version(torch.__version__, '1.10.0')
def make_anchors(feats, strides, grid_cell_offset=0.5):
"""Generate anchors from features."""
anchor_points, stride_tensor = [], []
assert feats is not None
dtype, device = feats[0].dtype, feats[0].device
for i, stride in enumerate(strides):
_, _, h, w = feats[i].shape
sx = torch.arange(end=w, device=device, dtype=dtype) + grid_cell_offset # shift x
sy = torch.arange(end=h, device=device, dtype=dtype) + grid_cell_offset # shift y
sy, sx = torch.meshgrid(sy, sx, indexing='ij') if TORCH_1_10 else torch.meshgrid(sy, sx)
anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2))
stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))
return torch.cat(anchor_points), torch.cat(stride_tensor)
def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
"""Transform distance(ltrb) to box(xywh or xyxy)."""
# 左上右下
lt, rb = torch.split(distance, 2, dim)
x1y1 = anchor_points - lt
x2y2 = anchor_points + rb
if xywh:
c_xy = (x1y1 + x2y2) / 2
wh = x2y2 - x1y1
return torch.cat((c_xy, wh), dim) # xywh bbox
return torch.cat((x1y1, x2y2), dim) # xyxy bbox
class DecodeBox():
def __init__(self, num_classes, input_shape):
super(DecodeBox, self).__init__()
self.num_classes = num_classes
self.bbox_attrs = 4 + num_classes
self.input_shape = input_shape
def decode_box(self, inputs):
# dbox batch_size, 4, 8400
# cls batch_size, 20, 8400
dbox, cls, origin_cls, anchors, strides = inputs
# 获得中心宽高坐标
dbox = dist2bbox(dbox, anchors.unsqueeze(0), xywh=True, dim=1) * strides
y = torch.cat((dbox, cls.sigmoid()), 1).permute(0, 2, 1)
# 进行归一化，到0~1之间
y[:, :, :4] = y[:, :, :4] / torch.Tensor([self.input_shape[1], self.input_shape[0], self.input_shape[1], self.input_shape[0]]).to(y.device)
return y
def decode_onnx_box(self, outputs):
"""
专门用于ONNX模型输出的解码方法
支持简化版（只输出box和cls）
"""
try:
if isinstance(outputs, list):
if len(outputs) == 2:
box_output = outputs[0] if isinstance(outputs[0], torch.Tensor) else torch.tensor(outputs[0])
cls_output = outputs[1] if isinstance(outputs[1], torch.Tensor) else torch.tensor(outputs[1])
# 获取形状信息
batch_size, _, num_anchors = box_output.shape
#print(f"Box形状: {box_output.shape}, Cls形状: {cls_output.shape}")
# 生成默认的anchor points和strides（针对YOLOv8 640x640输入）
strides = [8, 16, 32]
total_anchors = 0
anchor_points_list = []
stride_tensor_list = []
for stride in strides:
# 计算特征图尺寸
feature_h = self.input_shape[0] // stride
feature_w = self.input_shape[1] // stride
total_anchors += feature_h * feature_w
# 生成网格点
sx = torch.arange(0.5, feature_w + 0.5, dtype=torch.float32)
sy = torch.arange(0.5, feature_h + 0.5, dtype=torch.float32)
if TORCH_1_10:
sy_grid, sx_grid = torch.meshgrid(sy, sx, indexing='ij')
else:
sy_grid, sx_grid = torch.meshgrid(sy, sx)
# 展平并存储
anchors = torch.stack((sx_grid.flatten(), sy_grid.flatten()), dim=1)
anchor_points_list.append(anchors)
stride_tensor_list.append(torch.full((feature_h * feature_w, 1), stride, dtype=torch.float32))
# 合并所有anchor points
anchor_points = torch.cat(anchor_points_list, dim=0)
stride_tensor = torch.cat(stride_tensor_list, dim=0)
# 如果数量不匹配，进行调整
if anchor_points.shape[0] != num_anchors:
print(f"警告: anchor数量不匹配 (生成: {anchor_points.shape[0]}, 模型: {num_anchors})")
# 截断或重复以匹配
if anchor_points.shape[0] > num_anchors:
anchor_points = anchor_points[:num_anchors]
stride_tensor = stride_tensor[:num_anchors]
else:
repeat_times = num_anchors // anchor_points.shape[0]
remainder = num_anchors % anchor_points.shape[0]
if repeat_times > 0:
anchor_points = torch.cat([anchor_points] * repeat_times, dim=0)
stride_tensor = torch.cat([stride_tensor] * repeat_times, dim=0)
if remainder > 0:
anchor_points = torch.cat([anchor_points, anchor_points[:remainder]], dim=0)
stride_tensor = torch.cat([stride_tensor, stride_tensor[:remainder]], dim=0)
# 需要调整维度以匹配dist2bbox的期望输入
anchor_points = anchor_points.unsqueeze(0).permute(0, 2, 1)
stride_tensor = stride_tensor.unsqueeze(0).permute(0, 2, 1)
box_coords = dist2bbox(box_output, anchor_points, xywh=True, dim=1)
box_coords = box_coords * stride_tensor
# 应用sigmoid到分类输出
cls_scores = torch.sigmoid(cls_output)
# 合并box和cls输出 [batch, anchors, 4+num_classes]
y = torch.cat((box_coords.permute(0, 2, 1), cls_scores.permute(0, 2, 1)), dim=2)
else:
# 单个张量输出
y = outputs if isinstance(outputs, torch.Tensor) else torch.tensor(outputs)
# 归一化边界框坐标到0~1范围
if y.shape[-1] >= 4:
y[:, :, :4] = y[:, :, :4] / torch.tensor([
[self.input_shape[1], self.input_shape[0],
self.input_shape[1], self.input_shape[0]]
], dtype=torch.float32)
return y
except Exception as e:
print(f"ONNX解码出错: {e}")
import traceback
traceback.print_exc()
# 出错时返回原始输出
if isinstance(outputs, list):
return outputs[0] if len(outputs) > 0 else torch.empty(0)
return outputs
def yolo_correct_boxes(self, box_xy, box_wh, input_shape, image_shape, letterbox_image):
#-----------------------------------------------------------------#
# 把y轴放前面是因为方便预测框和图像的宽高进行相乘
#-----------------------------------------------------------------#
box_yx = box_xy[..., ::-1]
box_hw = box_wh[..., ::-1]
input_shape = np.array(input_shape)
image_shape = np.array(image_shape)
if letterbox_image:
#-----------------------------------------------------------------#
# 这里求出来的offset是图像有效区域相对于图像左上角的偏移情况
# new_shape指的是宽高缩放情况
#-----------------------------------------------------------------#
new_shape = np.round(image_shape * np.min(input_shape/image_shape))
offset = (input_shape - new_shape)/2./input_shape
scale = input_shape/new_shape
box_yx = (box_yx - offset) * scale
box_hw *= scale
box_mins = box_yx - (box_hw / 2.)
box_maxes = box_yx + (box_hw / 2.)
boxes = np.concatenate([box_mins[..., 0:1], box_mins[..., 1:2], box_maxes[..., 0:1], box_maxes[..., 1:2]], axis=-1)
boxes *= np.concatenate([image_shape, image_shape], axis=-1)
return boxes
def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):
#----------------------------------------------------------#
# 将预测结果的格式转换成左上角右下角的格式。
# prediction [batch_size, num_anchors, 85]
#----------------------------------------------------------#
# 确保prediction是tensor类型
if not isinstance(prediction, torch.Tensor):
prediction = torch.tensor(prediction)
box_corner = prediction.new(prediction.shape)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
#----------------------------------------------------------#
# 对种类预测部分取max。
# class_conf [num_anchors, 1] 种类置信度
# class_pred [num_anchors, 1] 种类
#----------------------------------------------------------#
# 确保索引不越界
if image_pred.shape[1] < 4 + num_classes:
print(f"警告: 预测输出形状不正确 {image_pred.shape}")
continue
class_conf, class_pred = torch.max(image_pred[:, 4:4 + num_classes], 1, keepdim=True)
#----------------------------------------------------------#
# 利用置信度进行第一轮筛选
#----------------------------------------------------------#
conf_mask = (class_conf[:, 0] >= conf_thres).squeeze()
#----------------------------------------------------------#
# 根据置信度进行预测结果的筛选
#----------------------------------------------------------#
image_pred = image_pred[conf_mask]
class_conf = class_conf[conf_mask]
class_pred = class_pred[conf_mask]
if not image_pred.size(0):
continue
#-------------------------------------------------------------------------#
# detections [num_anchors, 6]
# 6的内容为：x1, y1, x2, y2, class_conf, class_pred
#-------------------------------------------------------------------------#
detections = torch.cat((image_pred[:, :4], class_conf.float(), class_pred.float()), 1)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = detections[:, -1].cpu().unique()
if prediction.is_cuda:
unique_labels = unique_labels.cuda()
detections = detections.cuda()
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = detections[detections[:, -1] == c]
#------------------------------------------#
# 使用官方自带的非极大抑制会速度更快一些！
# 筛选出一定区域内，属于同一种类得分最大的框
#------------------------------------------#
keep = nms(
detections_class[:, :4],
detections_class[:, 4],
nms_thres
)
max_detections = detections_class[keep]
# Add max detections to outputs
output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))
if output[i] is not None:
output[i] = output[i].cpu().numpy()
box_xy, box_wh = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]
output[i][:, :4] = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)
return output
if __name__ == "__main__":
import matplotlib.pyplot as plt
import numpy as np
#---------------------------------------------------#
# 将预测值的每个特征层调成真实值
#---------------------------------------------------#
def get_anchors_and_decode(input, input_shape, anchors, anchors_mask, num_classes):
#-----------------------------------------------#
# input batch_size, 3 * (4 + 1 + num_classes), 20, 20
#-----------------------------------------------#
batch_size = input.size(0)
input_height = input.size(2)
input_width = input.size(3)
#-----------------------------------------------#
# 输入为640x640时 input_shape = [640, 640] input_height = 20, input_width = 20
# 640 / 20 = 32
# stride_h = stride_w = 32
#-----------------------------------------------#
stride_h = input_shape[0] / input_height
stride_w = input_shape[1] / input_width
#-------------------------------------------------#
# 此时获得的scaled_anchors大小是相对于特征层的
# anchor_width, anchor_height / stride_h, stride_w
#-------------------------------------------------#
scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in anchors[anchors_mask[2]]]
#-----------------------------------------------#
# batch_size, 3 * (4 + 1 + num_classes), 20, 20 =>
# batch_size, 3, 5 + num_classes, 20, 20 =>
# batch_size, 3, 20, 20, 4 + 1 + num_classes
#-----------------------------------------------#
prediction = input.view(batch_size, len(anchors_mask[2]),
num_classes + 5, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()
#-----------------------------------------------#
# 先验框的中心位置的调整参数
#-----------------------------------------------#
x = torch.sigmoid(prediction[..., 0])
y = torch.sigmoid(prediction[..., 1])
#-----------------------------------------------#
# 先验框的宽高调整参数
#-----------------------------------------------#
w = torch.sigmoid(prediction[..., 2])
h = torch.sigmoid(prediction[..., 3])
#-----------------------------------------------#
# 获得置信度，是否有物体 0 - 1
#-----------------------------------------------#
conf = torch.sigmoid(prediction[..., 4])
#-----------------------------------------------#
# 种类置信度 0 - 1
#-----------------------------------------------#
pred_cls = torch.sigmoid(prediction[..., 5:])
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
#----------------------------------------------------------#
# 生成网格，先验框中心，网格左上角
# batch_size,3,20,20
# range(20)
# [
# [0, 1, 2, 3 ……, 19],
# [0, 1, 2, 3 ……, 19],
# …… （20次）
# [0, 1, 2, 3 ……, 19]
# ] * (batch_size * 3)
# [batch_size, 3, 20, 20]
#
# [
# [0, 1, 2, 3 ……, 19],
# [0, 1, 2, 3 ……, 19],
# …… （20次）
# [0, 1, 2, 3 ……, 19]
# ].T * (batch_size * 3)
# [batch_size, 3, 20, 20]
#----------------------------------------------------------#
grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
batch_size * len(anchors_mask[2]), 1, 1).view(x.shape).type(FloatTensor)
grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(
batch_size * len(anchors_mask[2]), 1, 1).view(y.shape).type(FloatTensor)
#----------------------------------------------------------#
# 按照网格格式生成先验框的宽高
# batch_size, 3, 20 * 20 => batch_size, 3, 20, 20
# batch_size, 3, 20 * 20 => batch_size, 3, 20, 20
#----------------------------------------------------------#
anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)
#----------------------------------------------------------#
# 利用预测结果对先验框进行调整
# 首先调整先验框的中心，从先验框中心向右下角偏移
# 再调整先验框的宽高。
# x 0 ~ 1 => 0 ~ 2 => -0.5 ~ 1.5 + grid_x
# y 0 ~ 1 => 0 ~ 2 => -0.5 ~ 1.5 + grid_y
# w 0 ~ 1 => 0 ~ 2 => 0 ~ 4 * anchor_w
# h 0 ~ 1 => 0 ~ 2 => 0 ~ 4 * anchor_h
#----------------------------------------------------------#
pred_boxes = FloatTensor(prediction[..., :4].shape)
pred_boxes[..., 0] = x.data * 2. - 0.5 + grid_x
pred_boxes[..., 1] = y.data * 2. - 0.5 + grid_y
pred_boxes[..., 2] = (w.data * 2) ** 2 * anchor_w
pred_boxes[..., 3] = (h.data * 2) ** 2 * anchor_h
point_h = 5
point_w = 5
box_xy = pred_boxes[..., 0:2].cpu().numpy() * 32
box_wh = pred_boxes[..., 2:4].cpu().numpy() * 32
grid_x = grid_x.cpu().numpy() * 32
grid_y = grid_y.cpu().numpy() * 32
anchor_w = anchor_w.cpu().numpy() * 32
anchor_h = anchor_h.cpu().numpy() * 32
fig = plt.figure()
ax = fig.add_subplot(121)
from PIL import Image
img = Image.open("img/street.jpg").resize([640, 640])
plt.imshow(img, alpha=0.5)
plt.ylim(-30, 650)
plt.xlim(-30, 650)
plt.scatter(grid_x, grid_y)
plt.scatter(point_h * 32, point_w * 32, c='black')
plt.gca().invert_yaxis()
anchor_left = grid_x - anchor_w / 2
anchor_top = grid_y - anchor_h / 2
rect1 = plt.Rectangle([anchor_left[0, 0, point_h, point_w],anchor_top[0, 0, point_h, point_w]], \
anchor_w[0, 0, point_h, point_w],anchor_h[0, 0, point_h, point_w],color="r",fill=False)
rect2 = plt.Rectangle([anchor_left[0, 1, point_h, point_w],anchor_top[0, 1, point_h, point_w]], \
anchor_w[0, 1, point_h, point_w],anchor_h[0, 1, point_h, point_w],color="r",fill=False)
rect3 = plt.Rectangle([anchor_left[0, 2, point_h, point_w],anchor_top[0, 2, point_h, point_w]], \
anchor_w[0, 2, point_h, point_w],anchor_h[0, 2, point_h, point_w],color="r",fill=False)
ax.add_patch(rect1)
ax.add_patch(rect2)
ax.add_patch(rect3)
ax = fig.add_subplot(122)
plt.imshow(img, alpha=0.5)
plt.ylim(-30, 650)
plt.xlim(-30, 650)
plt.scatter(grid_x, grid_y)
plt.scatter(point_h * 32, point_w * 32, c='black')
plt.scatter(box_xy[0, :, point_h, point_w, 0], box_xy[0, :, point_h, point_w, 1], c='r')
plt.gca().invert_yaxis()
pre_left = box_xy[...,0] - box_wh[...,0] / 2
pre_top = box_xy[...,1] - box_wh[...,1] / 2
rect1 = plt.Rectangle([pre_left[0, 0, point_h, point_w], pre_top[0, 0, point_h, point_w]],\
box_wh[0, 0, point_h, point_w,0], box_wh[0, 0, point_h, point_w,1],color="r",fill=False)
rect2 = plt.Rectangle([pre_left[0, 1, point_h, point_w], pre_top[0, 1, point_h, point_w]],\
box_wh[0, 1, point_h, point_w,0], box_wh[0, 1, point_h, point_w,1],color="r",fill=False)
rect3 = plt.Rectangle([pre_left[0, 2, point_h, point_w], pre_top[0, 2, point_h, point_w]],\
box_wh[0, 2, point_h, point_w,0], box_wh[0, 2, point_h, point_w,1],color="r",fill=False)
ax.add_patch(rect1)
ax.add_patch(rect2)
ax.add_patch(rect3)
plt.show()
#
feat = torch.from_numpy(np.random.normal(0.2, 0.5, [4, 255, 20, 20])).float()
anchors = np.array([[116, 90], [156, 198], [373, 326], [30,61], [62,45], [59,119], [10,13], [16,30], [33,23]])
anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
get_anchors_and_decode(feat, [640, 640], anchors, anchors_mask, 80)3.1、初始化
3.1.1、初始化模型参数
3.1.2 初始化解码器
这份代码在utils_bbox文件中，主要功能是将检测到的数据进行解码，说人话就是先检测返回的结果是否正确，如果正确的话就转换成可以被opencv直接绘制检测框的数据返回。
3.2 模型检测
调用detect函数做检测：
def detect(self, image, conf_thres=0.5, nms_thres=0.3):
"""
对图像进行目标检测
Args:
image (PIL.Image): 输入图像
conf_thres (float): 置信度阈值
nms_thres (float): NMS阈值
Returns:
list: 检测结果列表，每个元素为[x0, y0, x1, y1, score, class_id]
"""
# 图像预处理
images, image_shape, original_image = self.preprocess_image(image)
# ONNX模型推理
try:
outputs = self.session.run(None, {self.input_name: images.cpu().numpy()})
outputs = [torch.from_numpy(output) for output in outputs]
except Exception as e:
print(f"ONNX模型推理错误: {e}")
return []
# 解码检测结果
try:
results = self.decodebox.decode_onnx_box(outputs)
except Exception as e:
print(f"解码错误: {e}")
return []
# 应用非极大值抑制
try:
results = self.decodebox.non_max_suppression(
results,
self.num_classes,
input_shape=self.input_shape,
image_shape=image_shape,
letterbox_image=True,
conf_thres=conf_thres,
nms_thres=nms_thres
)
except Exception as e:
print(f"NMS处理错误: {e}")
return []
# 处理检测结果
detections = []
if results[0] is not None:
top_label = np.array(results[0][:, 5], dtype='int32')
top_conf = results[0][:, 4]
top_boxes = results[0][:, :4]
for i in range(len(top_label)):
if top_conf[i] < conf_thres:
continue
box = top_boxes[i]
top, left, bottom, right = box
# 确保坐标在有效范围内
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(np.shape(original_image)[0], np.floor(bottom).astype('int32'))
right = min(np.shape(original_image)[1], np.floor(right).astype('int32'))
# 确保坐标顺序正确
x0 = min(left, right)
x1 = max(left, right)
y0 = min(top, bottom)
y1 = max(top, bottom)
if x1 > x0 and y1 > y0:
detections.append([x0, y0, x1, y1, top_conf[i], top_label[i]])
return detections
3.2.1 对图像做预处理
def preprocess_image(self, image):
"""
图像预处理
Args:
image (PIL.Image): 输入图像
Returns:
tuple: (处理后的图像tensor, 原始图像尺寸, 原始图像)
"""
image_shape = np.array(np.shape(image)[0:2])
# 转换为RGB格式
image = image.convert('RGB')
# 保存原始尺寸
original_width, original_height = image.size
# 调整图像大小（保持宽高比的letterbox）
scale = min(self.input_shape[0] / original_height, self.input_shape[1] / original_width)
new_width = int(original_width * scale)
new_height = int(original_height * scale)
# 创建新的图像并粘贴调整后的图像
resized_img = image.resize((new_width, new_height), Image.BICUBIC)
new_image = Image.new('RGB', self.input_shape, (128, 128, 128))
new_image.paste(resized_img, ((self.input_shape[1] - new_width) // 2, (self.input_shape[0] - new_height) // 2))
# 转换为numpy数组并归一化
image_data = np.array(new_image, dtype='float32') / 255.0
# 转换为CHW格式
image_data = np.transpose(image_data, (2, 0, 1))
# 添加批次维度
image_data = np.expand_dims(image_data, axis=0)
# 转换为tensor
images = torch.from_numpy(image_data)
return images, image_shape, image
3.2.2 模型推理
3.2.3 推理结果解码
调用的是uitls_bbox.py中DecodeBox类的decode_onnx_box函数，这个函数的功能其实就是，对检测到的数据进行检查，然后转化为可以方便被处理的数据。
def decode_onnx_box(self, outputs):
"""
专门用于ONNX模型输出的解码方法
支持简化版（只输出box和cls）
"""
try:
if isinstance(outputs, list):
if len(outputs) == 2:
box_output = outputs[0] if isinstance(outputs[0], torch.Tensor) else torch.tensor(outputs[0])
cls_output = outputs[1] if isinstance(outputs[1], torch.Tensor) else torch.tensor(outputs[1])
# 获取形状信息
batch_size, _, num_anchors = box_output.shape
#print(f"Box形状: {box_output.shape}, Cls形状: {cls_output.shape}")
# 生成默认的anchor points和strides（针对YOLOv8 640x640输入）
strides = [8, 16, 32]
total_anchors = 0
anchor_points_list = []
stride_tensor_list = []
for stride in strides:
# 计算特征图尺寸
feature_h = self.input_shape[0] // stride
feature_w = self.input_shape[1] // stride
total_anchors += feature_h * feature_w
# 生成网格点
sx = torch.arange(0.5, feature_w + 0.5, dtype=torch.float32)
sy = torch.arange(0.5, feature_h + 0.5, dtype=torch.float32)
if TORCH_1_10:
sy_grid, sx_grid = torch.meshgrid(sy, sx, indexing='ij')
else:
sy_grid, sx_grid = torch.meshgrid(sy, sx)
# 展平并存储
anchors = torch.stack((sx_grid.flatten(), sy_grid.flatten()), dim=1)
anchor_points_list.append(anchors)
stride_tensor_list.append(torch.full((feature_h * feature_w, 1), stride, dtype=torch.float32))
# 合并所有anchor points
anchor_points = torch.cat(anchor_points_list, dim=0)
stride_tensor = torch.cat(stride_tensor_list, dim=0)
# 如果数量不匹配，进行调整
if anchor_points.shape[0] != num_anchors:
print(f"警告: anchor数量不匹配 (生成: {anchor_points.shape[0]}, 模型: {num_anchors})")
# 截断或重复以匹配
if anchor_points.shape[0] > num_anchors:
anchor_points = anchor_points[:num_anchors]
stride_tensor = stride_tensor[:num_anchors]
else:
repeat_times = num_anchors // anchor_points.shape[0]
remainder = num_anchors % anchor_points.shape[0]
if repeat_times > 0:
anchor_points = torch.cat([anchor_points] * repeat_times, dim=0)
stride_tensor = torch.cat([stride_tensor] * repeat_times, dim=0)
if remainder > 0:
anchor_points = torch.cat([anchor_points, anchor_points[:remainder]], dim=0)
stride_tensor = torch.cat([stride_tensor, stride_tensor[:remainder]], dim=0)
# 需要调整维度以匹配dist2bbox的期望输入
anchor_points = anchor_points.unsqueeze(0).permute(0, 2, 1)
stride_tensor = stride_tensor.unsqueeze(0).permute(0, 2, 1)
box_coords = dist2bbox(box_output, anchor_points, xywh=True, dim=1)
box_coords = box_coords * stride_tensor
# 应用sigmoid到分类输出
cls_scores = torch.sigmoid(cls_output)
# 合并box和cls输出 [batch, anchors, 4+num_classes]
y = torch.cat((box_coords.permute(0, 2, 1), cls_scores.permute(0, 2, 1)), dim=2)
else:
# 单个张量输出
y = outputs if isinstance(outputs, torch.Tensor) else torch.tensor(outputs)
# 归一化边界框坐标到0~1范围
if y.shape[-1] >= 4:
y[:, :, :4] = y[:, :, :4] / torch.tensor([
[self.input_shape[1], self.input_shape[0],
self.input_shape[1], self.input_shape[0]]
], dtype=torch.float32)
return y
except Exception as e:
print(f"ONNX解码出错: {e}")
import traceback
traceback.print_exc()
# 出错时返回原始输出
if isinstance(outputs, list):
return outputs[0] if len(outputs) > 0 else torch.empty(0)
return outputs3.2.4 应用非极大值抑制
这一步的作用是，过滤掉一些低置信度的检测框，和避免同一个目标出现多个检测框
def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):
#----------------------------------------------------------#
# 将预测结果的格式转换成左上角右下角的格式。
# prediction [batch_size, num_anchors, 85]
#----------------------------------------------------------#
# 确保prediction是tensor类型
if not isinstance(prediction, torch.Tensor):
prediction = torch.tensor(prediction)
box_corner = prediction.new(prediction.shape)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
#----------------------------------------------------------#
# 对种类预测部分取max。
# class_conf [num_anchors, 1] 种类置信度
# class_pred [num_anchors, 1] 种类
#----------------------------------------------------------#
# 确保索引不越界
if image_pred.shape[1] < 4 + num_classes:
print(f"警告: 预测输出形状不正确 {image_pred.shape}")
continue
class_conf, class_pred = torch.max(image_pred[:, 4:4 + num_classes], 1, keepdim=True)
#----------------------------------------------------------#
# 利用置信度进行第一轮筛选
#----------------------------------------------------------#
conf_mask = (class_conf[:, 0] >= conf_thres).squeeze()
#----------------------------------------------------------#
# 根据置信度进行预测结果的筛选
#----------------------------------------------------------#
image_pred = image_pred[conf_mask]
class_conf = class_conf[conf_mask]
class_pred = class_pred[conf_mask]
if not image_pred.size(0):
continue
#-------------------------------------------------------------------------#
# detections [num_anchors, 6]
# 6的内容为：x1, y1, x2, y2, class_conf, class_pred
#-------------------------------------------------------------------------#
detections = torch.cat((image_pred[:, :4], class_conf.float(), class_pred.float()), 1)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = detections[:, -1].cpu().unique()
if prediction.is_cuda:
unique_labels = unique_labels.cuda()
detections = detections.cuda()
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = detections[detections[:, -1] == c]
#------------------------------------------#
# 使用官方自带的非极大抑制会速度更快一些！
# 筛选出一定区域内，属于同一种类得分最大的框
#------------------------------------------#
keep = nms(
detections_class[:, :4],
detections_class[:, 4],
nms_thres
)
max_detections = detections_class[keep]
# Add max detections to outputs
output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))
if output[i] is not None:
output[i] = output[i].cpu().numpy()
box_xy, box_wh = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]
output[i][:, :4] = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)
return output
基本上到这一步就差不多了，接下来就是对数据再进行一轮检查，检查是否再有限的范围内，然后将检测框绘制出来。
3.2.5 检查并绘制
检查检测框坐标正确后，进行绘制。
def draw_detections(self, image, detections, class_names=None):
"""
在图像上绘制检测结果
Args:
image (PIL.Image): 原始图像
detections (list): 检测结果列表
class_names (list): 类别名称列表
Returns:
PIL.Image: 绘制了检测结果的图像
"""
from PIL import ImageDraw, ImageFont
import numpy as np
# 复制原图用于绘制结果
result_image = image.copy()
draw = ImageDraw.Draw(result_image)
# 设置字体
try:
font_size = max(16, int(np.floor(3e-2 * np.shape(image)[1] + 10)))
font = ImageFont.truetype(font='model_data/simhei.ttf', size=font_size)
except:
try:
font = ImageFont.truetype(font='arial.ttf', size=font_size)
except:
font = ImageFont.load_default()
thickness = max(2, int(max((np.shape(image)[0] + np.shape(image)[1]) // self.input_shape[0], 1)))
# 绘制检测框
for i, det in enumerate(detections):
if len(det) < 6:
continue
x0, y0, x1, y1, score, class_id = det[:6]
# 使用类别名称
predicted_class = class_names[int(class_id)] if class_names and int(class_id) < len(class_names) else f"Class {int(class_id)}"
# 绘制边界框
color = (255, 0, 0)
x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
if x1 > x0 and y1 > y0:
draw.rectangle([x0, y0, x1, y1], outline=color, width=thickness)
label = f'{predicted_class} {score:.2f}'
draw.text((x0, y0), label, fill=color, font=font)
return result_image
到这里整个流程就都结束了！！！！
最后，由ONNX转OpenVINO格式的方法，其实和本文的方法是相通的，我在本文就不做分析了，这里推荐一个文章，他的方法比较直接：onnx模型转openvinoIR模型+推理（学习笔记）_onnx转openvino-CSDN博客
智能硬件社区聚焦AI智能硬件技术生态，汇聚嵌入式AI、物联网硬件开发者，打造交流分享平台，同步全国赛事资讯、开展 OPC 核心人才招募，助力技术落地与开发者成长。
更多推荐
36
0- 0
已为社区贡献2条内容
所有评论(0)