1.osnet1.1 源码model_zoo:https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.htmlosnet源码https://github.com/KaiyangZhou/deep-person-reid1.2 分为三步冻结部分层进行微调1.stage1使用torchreid的默认配置冻结全部网络只解冻训练classifier BNNeck。2.stage2使用torchreid的默认配置冻结全部网络只解冻 classifier BNNeck解冻高层 backboneconv53.stage3使用torchreid的默认配置解冻所有层1.3 自动化三阶段OSNet模型训练脚本 train_osnet_3stage_auto.py将以下代码放至./model/deep-person-reid下的train_osnet_3stage_auto.py自动化三阶段OSNet模型训练脚本 train_osnet_3stage_auto.py 此脚本会自动依次执行Stage 1 - Stage 2 - Stage 3的训练无需手动干预 import torch import torchreid from torchreid.utils import set_random_seed import os import glob # # 0. 基础配置 # set_random_seed(42) # 数据集配置 DATASET_NAME ten_object # 数据集 pubfigface DATASET_ROOT /ReID/datasets/ # 数据集根路径数据集路径在以下注册所以这里只需要根路径别搞错了 BASE_SAVE_DIR ./log/osnet_ten_object_auto1 # 基础保存目录 MODEL_NAME osnet_x1_0 # 选用预训练模型一般不用改 IMG_HEIGHT 256 IMG_WIDTH 128 BATCH_SIZE 64 DEVICE cuda if torch.cuda.is_available() else cpu # 初始化数据管理器 print(正在初始化数据管理器...) datamanager torchreid.data.ImageDataManager( rootDATASET_ROOT, sourcesDATASET_NAME, targetsDATASET_NAME, heightIMG_HEIGHT, widthIMG_WIDTH, batch_size_trainBATCH_SIZE, batch_size_testBATCH_SIZE, transforms[random_flip, random_crop], num_instances4 ) def find_latest_model(stage_num, base_dir): 查找指定阶段的最新模型文件 # 实际保存路径就是 base_dir/model/不需要再加 stage 后缀 model_dir os.path.join(base_dir, model) pattern os.path.join(model_dir, model.pth.tar-*) model_files glob.glob(pattern) if not model_files: print(f在 {model_dir} 中未找到模型文件) return None # 按epoch数字排序返回最大epoch的模型 def extract_epoch(path): try: # 提取文件名中的epoch数字 basename os.path.basename(path) # 从文件名中提取epoch数字例如从 model.pth.tar-40 提取 40 parts basename.split(-) if len(parts) 0: epoch_part parts[-1].replace(.pth.tar, ) epoch int(epoch_part) return epoch except: return 0 latest_model max(model_files, keyextract_epoch) return latest_model def get_stage_config(stage_num, prev_model_pathNone): 获取指定阶段的配置 epochs_map {1: 40, 2: 80, 3: 40} config { save_dir: f{BASE_SAVE_DIR}_stage{stage_num}, epochs: epochs_map.get(stage_num, 40), lr_backbone: 0.0, lr_head: 1e-3 } if stage_num 1: config[lr_backbone] 0.0 config[lr_head] 1e-3 elif stage_num 2: config[lr_backbone] 1e-5 config[lr_head] 1e-4 elif stage_num 3: config[lr_backbone] 3e-4 config[lr_head] 3e-4 return config def build_and_load_model(stage_num, prev_model_pathNone): 构建模型并加载相应阶段的权重 print(f构建阶段 {stage_num} 的模型...) model torchreid.models.build_model( nameMODEL_NAME, num_classesdatamanager.num_train_pids, losssoftmax, pretrainedFalse # 关闭自动下载 ) model model.to(DEVICE) # 确定要加载的权重 if stage_num 1: # 第一阶段加载ImageNet预训练权重 local_weights ../osnet_x1_0_imagenet.pth print(f加载基础预训练权重: {local_weights}) else: # 后续阶段加载前一阶段的模型权重 if prev_model_path and os.path.exists(prev_model_path): local_weights prev_model_path print(f加载前一阶段模型权重: {local_weights}) else: raise FileNotFoundError(f找不到第 {stage_num-1} 阶段的模型文件: {prev_model_path}) # 加载权重 state_dict torch.load(local_weights, map_locationDEVICE, weights_onlyFalse) # 兼容 torchreid / checkpoint 两种格式 if state_dict in state_dict: state_dict state_dict[state_dict] # 删除 classifier 权重对于第1阶段避免 size mismatch if stage_num 1: for key in list(state_dict.keys()): if key.startswith(classifier): del state_dict[key] print(移除了预训练模型中的分类头权重) missing, unexpected model.load_state_dict(state_dict, strictFalse) print(加载模型权重成功) print(缺失的键:, missing) print(意外的键:, unexpected) return model def setup_training_parameters(model, stage_num): 设置训练参数冻结/解冻策略 print(f设置阶段 {stage_num} 的训练参数...) # 先全部冻结 for param in model.parameters(): param.requires_grad False # Stage 1: 只训练 classifier BNNeck if stage_num 1: for name, param in model.named_parameters(): if classifier in name or bn in name: param.requires_grad True # 必须重新初始化分类头 model.classifier.reset_parameters() lr_backbone 0.0 lr_head 1e-3 # Stage 2: 解冻 classifier BNNeck 高层 backboneconv5 elif stage_num 2: for name, param in model.named_parameters(): if classifier in name or bn in name: param.requires_grad True # 解冻高层 backboneconv5 for name, param in model.named_parameters(): if conv5 in name: param.requires_grad True lr_backbone 1e-5 lr_head 1e-4 # Stage 3: 全模型训练 elif stage_num 3: for param in model.parameters(): param.requires_grad True lr_backbone 3e-4 lr_head 3e-4 else: raise ValueError(STAGE must be 1, 2, or 3) # 设置优化器参数 params list(model.named_parameters()) backbone_params [p for n, p in params if classifier not in n] head_params [p for n, p in params if classifier in n] optimizer torch.optim.Adam( [ {params: backbone_params, lr: lr_backbone}, {params: head_params, lr: lr_head}, ], weight_decay5e-4 ) # 学习率调度 scheduler torch.optim.lr_scheduler.StepLR( optimizer, step_size20, gamma0.1 ) return optimizer, scheduler def train_stage(stage_num, prev_model_pathNone): 训练单个阶段 print(f\n{*50}) print(f开始第 {stage_num} 阶段训练) print(f{*50}) # 构建模型并加载权重 model build_and_load_model(stage_num, prev_model_path) # 设置训练参数 optimizer, scheduler setup_training_parameters(model, stage_num) # 获取阶段配置 config get_stage_config(stage_num) # 创建Trainer engine torchreid.engine.ImageSoftmaxEngine( datamanager, model, optimizeroptimizer, schedulerscheduler, label_smoothTrue ) # 开始训练 print(f开始训练保存到: {config[save_dir]}) engine.run( save_dirconfig[save_dir], max_epochconfig[epochs], eval_freq5, print_freq20 ) # 等待一段时间让文件系统完成写入 import time print(等待模型文件写入完成...) time.sleep(5) # 等待5秒让文件系统完成写入 # 查找并返回此阶段训练完成的模型路径 latest_model find_latest_model(stage_num, config[save_dir]) if latest_model: print(f阶段 {stage_num} 训练完成最新模型: {latest_model}) else: print(f警告: 未找到阶段 {stage_num} 的模型文件) # 再次尝试查找以防文件还在写入过程中 time.sleep(10) latest_model find_latest_model(stage_num, config[save_dir]) if latest_model: print(f阶段 {stage_num} 训练完成最新模型: {latest_model}) else: print(f再次检查后仍未找到阶段 {stage_num} 的模型文件) return latest_model def main(): 主函数执行三阶段训练 print(开始自动化三阶段训练流程) print(*60) prev_model_path None # 依次执行三个阶段 for stage in range(1, 4): prev_model_path train_stage(stage, prev_model_path) if not prev_model_path: print(f错误: 阶段 {stage} 训练失败无法继续后续阶段) return print(\n *60) print(所有三个阶段训练完成) print(f最终模型保存在: {BASE_SAVE_DIR}_stage3) print(训练流程结束) if __name__ __main__: main()1.4 如若自制数据集格式如下命名规则如0001_c001_00016450_0.jpg,第一个下划线前是id号0001第二个下划线前是相机号c0001,然后后面是图片名字最后是扩展名.jpgID数应大于等于50否则效果可能会下降格式要求每个目标至少需要4张图片4张图片需要两个不同的相机号train中放训练集包含cam1和cam2的图片query是目标集用来验证用的gallery是搜索集也就是使用query中的图片在gallery中进行搜索注意query中的id和gallery中的id必须一样(不然搜不到)cam号必须不一样(算法要求)且必须与train中的id不一样(不然即参与训练又参与验证会导致map升高)# 示例数据集格式 aircraft_reid/ ├── train/ │ ├── 0001_c001_00016450.jpg │ ├── 0001_c002_00016915.jpg │ ├── 0001_c002_00014680.jpg | ├── query/ │ ├── 0002_c001_00030600.jpg │ └── gallery/ ├── 0002_c002_00030600.jpg │ ├── 0005_c002_00075750.jpg1.5 数据集划分可使用如下脚本进行yolo格式的数据集的划分划分前数据集格式和分类数据集一致每个文件夹为一类物体的所有图片有若干文件夹更改其中的DATA_YAML,OUTPUT_DIR,PREFIX 根据YOLO格式的标签从图片中裁剪出目标并按照指定前缀和类别名保存到对应文件夹中。 import os import yaml import cv2 # 配置 DATA_YAML classes3.yaml # 数据集的yaml文件路径,yaml文件为yolo格式如下 path: /data/VOCdevkit/ # 数据集目录 train: - train/images val: - valid/images test: - test/images # Classes names: 0: class0 1: class1 2: class2 OUTPUT_DIR classes3_output # 输出目录 PREFIX classes3 # 你指定的前缀例如 classes3 IMG_EXTS [.jpg, .jpeg, .png, .bmp] # def load_yaml(yaml_path): with open(yaml_path, r, encodingutf-8) as f: data yaml.safe_load(f) return data def is_image(file): return os.path.splitext(file)[1].lower() in IMG_EXTS def find_image_path(img_dir, base_name): for ext in IMG_EXTS: img_path os.path.join(img_dir, base_name ext) if os.path.exists(img_path): return img_path return None def ensure_dir(path): if not os.path.exists(path): os.makedirs(path) def yolo_to_xyxy(img_w, img_h, x, y, w, h): YOLO格式 - 像素坐标 x1 int((x - w / 2) * img_w) y1 int((y - h / 2) * img_h) x2 int((x w / 2) * img_w) y2 int((y h / 2) * img_h) # 边界裁剪 x1 max(0, x1) y1 max(0, y1) x2 min(img_w - 1, x2) y2 min(img_h - 1, y2) return x1, y1, x2, y2 def get_all_image_label_pairs(root_path): 通用版本 支持任意结构只要满足 images / labels 对应关系 pairs [] for dirpath, _, filenames in os.walk(root_path): # 关键路径中包含 labels if labels not in dirpath: continue # 推导对应的 images 路径 img_dir dirpath.replace(os.sep labels, os.sep images) if not os.path.exists(img_dir): # print(f⚠️ 没有对应 images: {img_dir}) continue for file in filenames: if not file.endswith(.txt): continue base os.path.splitext(file)[0] label_path os.path.join(dirpath, file) img_path find_image_path(img_dir, base) if img_path: pairs.append((img_path, label_path)) else: print(f⚠️ 图片不存在: {base}) return pairs def main(): data load_yaml(DATA_YAML) root_path data[path] names data[names] # {0: M1A1, 1: T72, ...} print(f数据集路径: {root_path}) print(f类别: {names}) pairs get_all_image_label_pairs(root_path) print(f共找到 {len(pairs)} 对图片标签) # 每个类别计数 counters {int(k): 0 for k in names.keys()} for img_path, label_path in pairs: img cv2.imread(img_path) if img is None: continue h, w img.shape[:2] with open(label_path, r) as f: lines f.readlines() if len(lines) 0: continue for line in lines: parts line.strip().split() if len(parts) ! 5: continue cls_id int(parts[0]) x, y, bw, bh map(float, parts[1:]) # 新增类别过滤 if cls_id not in names: print(f⚠️ 跳过未知类别 {cls_id} - {label_path}) continue x1, y1, x2, y2 yolo_to_xyxy(w, h, x, y, bw, bh) crop img[y1:y2, x1:x2] if crop.size 0: continue cls_name names[cls_id] # 输出文件夹 out_dir os.path.join(OUTPUT_DIR, f{PREFIX}_{cls_name}) ensure_dir(out_dir) # 文件名 idx counters[cls_id] save_path os.path.join(out_dir, f{PREFIX}_{cls_name}_{idx:06d}.jpg) cv2.imwrite(save_path, crop) counters[cls_id] 1 print(处理完成) if __name__ __main__: main()1.6 制作osnet数据集在1.4生成的众多文件夹内同目录下创建以下脚本制作osnet数据集,更改其中的SRC_DIR,OUT_DIR# 所有种类是同种类时 import os import shutil import random # 配置 SRC_DIR classes3_output OUT_DIR reid_dataset_classes3 TRAIN_RATIO 0.7 # 70% 类别用于训练 random.seed(0) # def ensure_dir(path): if not os.path.exists(path): os.makedirs(path) def main(): classes [d for d in os.listdir(SRC_DIR) if os.path.isdir(os.path.join(SRC_DIR, d))] classes.sort() print(f总类别数: {len(classes)}) random.shuffle(classes) split_idx int(len(classes) * TRAIN_RATIO) train_classes classes[:split_idx] test_classes classes[split_idx:] # 创建输出目录 train_dir os.path.join(OUT_DIR, train) query_dir os.path.join(OUT_DIR, query) gallery_dir os.path.join(OUT_DIR, gallery) ensure_dir(train_dir) ensure_dir(query_dir) ensure_dir(gallery_dir) # 处理 train 类 train_id 1 for cls in train_classes: cls_path os.path.join(SRC_DIR, cls) imgs [f for f in os.listdir(cls_path) if f.endswith(.jpg)] for i, img_name in enumerate(imgs): cam_id 1 if i % 2 0 else 2 new_name f{train_id:04d}_c{cam_id:03d}_{i:08d}.jpg shutil.copy(os.path.join(cls_path, img_name), os.path.join(train_dir, new_name)) train_id 1 # 处理 query/gallery 类 test_id 1001 for cls in test_classes: cls_path os.path.join(SRC_DIR, cls) imgs [f for f in os.listdir(cls_path) if f.endswith(.jpg)] random.shuffle(imgs) # 至少两张用于 query/gallery query_img imgs[0] gallery_img imgs[1] shutil.copy(os.path.join(cls_path, query_img), os.path.join(query_dir, f{test_id:04d}_c001_00000000.jpg)) shutil.copy(os.path.join(cls_path, gallery_img), os.path.join(gallery_dir, f{test_id:04d}_c002_00000001.jpg)) # 剩余放 gallery for i, img_name in enumerate(imgs[2:]): new_name f{test_id:04d}_c002_{i2:08d}.jpg shutil.copy(os.path.join(cls_path, img_name), os.path.join(gallery_dir, new_name)) test_id 1 print(✅ 方案一数据集划分完成) if __name__ __main__: main()1.7 注册数据集osnet不支持像yolo那样以文件夹名为数据集需注册在deep-person-reid/torchreid/data/datasets/下新建如classes3.py文件内容如下更改其中的self.train_dir、self.query_dir、self.gallery_dir路径更改其中的dataset_dir更改其中的class classes3类名和对应的super(classes3, self).__init__(train, query, gallery, **kwargs)# classes3.py from __future__ import absolute_import import os.path as osp import glob import re from ..dataset import ImageDataset class classes3(ImageDataset): dataset_dir reid_dataset_classes3 # osnet的数据集路径 def __init__(self, root, **kwargs): self.root osp.abspath(osp.expanduser(root)) self.dataset_dir osp.join(self.root, self.dataset_dir) self.train_dir osp.join(self.dataset_dir, train) self.query_dir osp.join(self.dataset_dir, query) self.gallery_dir osp.join(self.dataset_dir, gallery) required_files [ self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir ] self.check_before_run(required_files) train self.process_dir(self.train_dir, relabelTrue) query self.process_dir(self.query_dir, relabelFalse) gallery self.process_dir(self.gallery_dir, relabelFalse) super(classes3, self).__init__(train, query, gallery, **kwargs) def process_dir(self, dir_path, relabelFalse): img_paths glob.glob(osp.join(dir_path, *.jpg)) pattern re.compile(r([-\d])_c(\d)) pid_container set() for img_path in img_paths: pid, _ map(int, pattern.search(img_path).groups()) if pid 0: continue pid_container.add(pid) pid2label {pid: label for label, pid in enumerate(pid_container)} data [] for img_path in img_paths: pid, camid map(int, pattern.search(img_path).groups()) if pid 0: continue camid - 1 # camid 从 0 开始 if relabel: pid pid2label[pid] data.append((img_path, pid, camid)) return data在 datasets/init.py 中注册文件路径model/deep-person-reid/torchreid/data/datasets/__init__.py添加:from .image.classes3 import classes3添加到 factory:找到__image_datasets {.....},在其中添加如下__image_datasets { ... classes3: classes3, }1.8 更改训练脚本train_osnet_3stage_auto.py完成后训练脚本的路径改为DATASET_NAME reid_dataset_classes3DATASET_ROOT改为classes3数据集的根路径2.yolo的冻结训练只需加一个参数freezefreeze 值含义22只训练 Detect Head极端小数据10Backbone 全冻结推荐起点15Backbone 部分 Neck0不冻结代码from ultralytics import YOLO import os os.environ[YOLO_DISABLE_AMP_CHECK] 1 model YOLO(./yolov8s.pt) # 加载模型 results model.train(datatruck.yaml, epochs50, workers16,imgsz640, batch64, device0,1,resumeFalse, freeze22, patience20)# cacheTrue # python -m torch.distributed.run --nproc_per_node2 ./train_detect.py
osnet模型和yolo模型的微调：冻结训练
发布时间：2026/7/29 19:38:03
1.osnet1.1 源码model_zoo:https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.htmlosnet源码https://github.com/KaiyangZhou/deep-person-reid1.2 分为三步冻结部分层进行微调1.stage1使用torchreid的默认配置冻结全部网络只解冻训练classifier BNNeck。2.stage2使用torchreid的默认配置冻结全部网络只解冻 classifier BNNeck解冻高层 backboneconv53.stage3使用torchreid的默认配置解冻所有层1.3 自动化三阶段OSNet模型训练脚本 train_osnet_3stage_auto.py将以下代码放至./model/deep-person-reid下的train_osnet_3stage_auto.py自动化三阶段OSNet模型训练脚本 train_osnet_3stage_auto.py 此脚本会自动依次执行Stage 1 - Stage 2 - Stage 3的训练无需手动干预 import torch import torchreid from torchreid.utils import set_random_seed import os import glob # # 0. 基础配置 # set_random_seed(42) # 数据集配置 DATASET_NAME ten_object # 数据集 pubfigface DATASET_ROOT /ReID/datasets/ # 数据集根路径数据集路径在以下注册所以这里只需要根路径别搞错了 BASE_SAVE_DIR ./log/osnet_ten_object_auto1 # 基础保存目录 MODEL_NAME osnet_x1_0 # 选用预训练模型一般不用改 IMG_HEIGHT 256 IMG_WIDTH 128 BATCH_SIZE 64 DEVICE cuda if torch.cuda.is_available() else cpu # 初始化数据管理器 print(正在初始化数据管理器...) datamanager torchreid.data.ImageDataManager( rootDATASET_ROOT, sourcesDATASET_NAME, targetsDATASET_NAME, heightIMG_HEIGHT, widthIMG_WIDTH, batch_size_trainBATCH_SIZE, batch_size_testBATCH_SIZE, transforms[random_flip, random_crop], num_instances4 ) def find_latest_model(stage_num, base_dir): 查找指定阶段的最新模型文件 # 实际保存路径就是 base_dir/model/不需要再加 stage 后缀 model_dir os.path.join(base_dir, model) pattern os.path.join(model_dir, model.pth.tar-*) model_files glob.glob(pattern) if not model_files: print(f在 {model_dir} 中未找到模型文件) return None # 按epoch数字排序返回最大epoch的模型 def extract_epoch(path): try: # 提取文件名中的epoch数字 basename os.path.basename(path) # 从文件名中提取epoch数字例如从 model.pth.tar-40 提取 40 parts basename.split(-) if len(parts) 0: epoch_part parts[-1].replace(.pth.tar, ) epoch int(epoch_part) return epoch except: return 0 latest_model max(model_files, keyextract_epoch) return latest_model def get_stage_config(stage_num, prev_model_pathNone): 获取指定阶段的配置 epochs_map {1: 40, 2: 80, 3: 40} config { save_dir: f{BASE_SAVE_DIR}_stage{stage_num}, epochs: epochs_map.get(stage_num, 40), lr_backbone: 0.0, lr_head: 1e-3 } if stage_num 1: config[lr_backbone] 0.0 config[lr_head] 1e-3 elif stage_num 2: config[lr_backbone] 1e-5 config[lr_head] 1e-4 elif stage_num 3: config[lr_backbone] 3e-4 config[lr_head] 3e-4 return config def build_and_load_model(stage_num, prev_model_pathNone): 构建模型并加载相应阶段的权重 print(f构建阶段 {stage_num} 的模型...) model torchreid.models.build_model( nameMODEL_NAME, num_classesdatamanager.num_train_pids, losssoftmax, pretrainedFalse # 关闭自动下载 ) model model.to(DEVICE) # 确定要加载的权重 if stage_num 1: # 第一阶段加载ImageNet预训练权重 local_weights ../osnet_x1_0_imagenet.pth print(f加载基础预训练权重: {local_weights}) else: # 后续阶段加载前一阶段的模型权重 if prev_model_path and os.path.exists(prev_model_path): local_weights prev_model_path print(f加载前一阶段模型权重: {local_weights}) else: raise FileNotFoundError(f找不到第 {stage_num-1} 阶段的模型文件: {prev_model_path}) # 加载权重 state_dict torch.load(local_weights, map_locationDEVICE, weights_onlyFalse) # 兼容 torchreid / checkpoint 两种格式 if state_dict in state_dict: state_dict state_dict[state_dict] # 删除 classifier 权重对于第1阶段避免 size mismatch if stage_num 1: for key in list(state_dict.keys()): if key.startswith(classifier): del state_dict[key] print(移除了预训练模型中的分类头权重) missing, unexpected model.load_state_dict(state_dict, strictFalse) print(加载模型权重成功) print(缺失的键:, missing) print(意外的键:, unexpected) return model def setup_training_parameters(model, stage_num): 设置训练参数冻结/解冻策略 print(f设置阶段 {stage_num} 的训练参数...) # 先全部冻结 for param in model.parameters(): param.requires_grad False # Stage 1: 只训练 classifier BNNeck if stage_num 1: for name, param in model.named_parameters(): if classifier in name or bn in name: param.requires_grad True # 必须重新初始化分类头 model.classifier.reset_parameters() lr_backbone 0.0 lr_head 1e-3 # Stage 2: 解冻 classifier BNNeck 高层 backboneconv5 elif stage_num 2: for name, param in model.named_parameters(): if classifier in name or bn in name: param.requires_grad True # 解冻高层 backboneconv5 for name, param in model.named_parameters(): if conv5 in name: param.requires_grad True lr_backbone 1e-5 lr_head 1e-4 # Stage 3: 全模型训练 elif stage_num 3: for param in model.parameters(): param.requires_grad True lr_backbone 3e-4 lr_head 3e-4 else: raise ValueError(STAGE must be 1, 2, or 3) # 设置优化器参数 params list(model.named_parameters()) backbone_params [p for n, p in params if classifier not in n] head_params [p for n, p in params if classifier in n] optimizer torch.optim.Adam( [ {params: backbone_params, lr: lr_backbone}, {params: head_params, lr: lr_head}, ], weight_decay5e-4 ) # 学习率调度 scheduler torch.optim.lr_scheduler.StepLR( optimizer, step_size20, gamma0.1 ) return optimizer, scheduler def train_stage(stage_num, prev_model_pathNone): 训练单个阶段 print(f\n{*50}) print(f开始第 {stage_num} 阶段训练) print(f{*50}) # 构建模型并加载权重 model build_and_load_model(stage_num, prev_model_path) # 设置训练参数 optimizer, scheduler setup_training_parameters(model, stage_num) # 获取阶段配置 config get_stage_config(stage_num) # 创建Trainer engine torchreid.engine.ImageSoftmaxEngine( datamanager, model, optimizeroptimizer, schedulerscheduler, label_smoothTrue ) # 开始训练 print(f开始训练保存到: {config[save_dir]}) engine.run( save_dirconfig[save_dir], max_epochconfig[epochs], eval_freq5, print_freq20 ) # 等待一段时间让文件系统完成写入 import time print(等待模型文件写入完成...) time.sleep(5) # 等待5秒让文件系统完成写入 # 查找并返回此阶段训练完成的模型路径 latest_model find_latest_model(stage_num, config[save_dir]) if latest_model: print(f阶段 {stage_num} 训练完成最新模型: {latest_model}) else: print(f警告: 未找到阶段 {stage_num} 的模型文件) # 再次尝试查找以防文件还在写入过程中 time.sleep(10) latest_model find_latest_model(stage_num, config[save_dir]) if latest_model: print(f阶段 {stage_num} 训练完成最新模型: {latest_model}) else: print(f再次检查后仍未找到阶段 {stage_num} 的模型文件) return latest_model def main(): 主函数执行三阶段训练 print(开始自动化三阶段训练流程) print(*60) prev_model_path None # 依次执行三个阶段 for stage in range(1, 4): prev_model_path train_stage(stage, prev_model_path) if not prev_model_path: print(f错误: 阶段 {stage} 训练失败无法继续后续阶段) return print(\n *60) print(所有三个阶段训练完成) print(f最终模型保存在: {BASE_SAVE_DIR}_stage3) print(训练流程结束) if __name__ __main__: main()1.4 如若自制数据集格式如下命名规则如0001_c001_00016450_0.jpg,第一个下划线前是id号0001第二个下划线前是相机号c0001,然后后面是图片名字最后是扩展名.jpgID数应大于等于50否则效果可能会下降格式要求每个目标至少需要4张图片4张图片需要两个不同的相机号train中放训练集包含cam1和cam2的图片query是目标集用来验证用的gallery是搜索集也就是使用query中的图片在gallery中进行搜索注意query中的id和gallery中的id必须一样(不然搜不到)cam号必须不一样(算法要求)且必须与train中的id不一样(不然即参与训练又参与验证会导致map升高)# 示例数据集格式 aircraft_reid/ ├── train/ │ ├── 0001_c001_00016450.jpg │ ├── 0001_c002_00016915.jpg │ ├── 0001_c002_00014680.jpg | ├── query/ │ ├── 0002_c001_00030600.jpg │ └── gallery/ ├── 0002_c002_00030600.jpg │ ├── 0005_c002_00075750.jpg1.5 数据集划分可使用如下脚本进行yolo格式的数据集的划分划分前数据集格式和分类数据集一致每个文件夹为一类物体的所有图片有若干文件夹更改其中的DATA_YAML,OUTPUT_DIR,PREFIX 根据YOLO格式的标签从图片中裁剪出目标并按照指定前缀和类别名保存到对应文件夹中。 import os import yaml import cv2 # 配置 DATA_YAML classes3.yaml # 数据集的yaml文件路径,yaml文件为yolo格式如下 path: /data/VOCdevkit/ # 数据集目录 train: - train/images val: - valid/images test: - test/images # Classes names: 0: class0 1: class1 2: class2 OUTPUT_DIR classes3_output # 输出目录 PREFIX classes3 # 你指定的前缀例如 classes3 IMG_EXTS [.jpg, .jpeg, .png, .bmp] # def load_yaml(yaml_path): with open(yaml_path, r, encodingutf-8) as f: data yaml.safe_load(f) return data def is_image(file): return os.path.splitext(file)[1].lower() in IMG_EXTS def find_image_path(img_dir, base_name): for ext in IMG_EXTS: img_path os.path.join(img_dir, base_name ext) if os.path.exists(img_path): return img_path return None def ensure_dir(path): if not os.path.exists(path): os.makedirs(path) def yolo_to_xyxy(img_w, img_h, x, y, w, h): YOLO格式 - 像素坐标 x1 int((x - w / 2) * img_w) y1 int((y - h / 2) * img_h) x2 int((x w / 2) * img_w) y2 int((y h / 2) * img_h) # 边界裁剪 x1 max(0, x1) y1 max(0, y1) x2 min(img_w - 1, x2) y2 min(img_h - 1, y2) return x1, y1, x2, y2 def get_all_image_label_pairs(root_path): 通用版本 支持任意结构只要满足 images / labels 对应关系 pairs [] for dirpath, _, filenames in os.walk(root_path): # 关键路径中包含 labels if labels not in dirpath: continue # 推导对应的 images 路径 img_dir dirpath.replace(os.sep labels, os.sep images) if not os.path.exists(img_dir): # print(f⚠️ 没有对应 images: {img_dir}) continue for file in filenames: if not file.endswith(.txt): continue base os.path.splitext(file)[0] label_path os.path.join(dirpath, file) img_path find_image_path(img_dir, base) if img_path: pairs.append((img_path, label_path)) else: print(f⚠️ 图片不存在: {base}) return pairs def main(): data load_yaml(DATA_YAML) root_path data[path] names data[names] # {0: M1A1, 1: T72, ...} print(f数据集路径: {root_path}) print(f类别: {names}) pairs get_all_image_label_pairs(root_path) print(f共找到 {len(pairs)} 对图片标签) # 每个类别计数 counters {int(k): 0 for k in names.keys()} for img_path, label_path in pairs: img cv2.imread(img_path) if img is None: continue h, w img.shape[:2] with open(label_path, r) as f: lines f.readlines() if len(lines) 0: continue for line in lines: parts line.strip().split() if len(parts) ! 5: continue cls_id int(parts[0]) x, y, bw, bh map(float, parts[1:]) # 新增类别过滤 if cls_id not in names: print(f⚠️ 跳过未知类别 {cls_id} - {label_path}) continue x1, y1, x2, y2 yolo_to_xyxy(w, h, x, y, bw, bh) crop img[y1:y2, x1:x2] if crop.size 0: continue cls_name names[cls_id] # 输出文件夹 out_dir os.path.join(OUTPUT_DIR, f{PREFIX}_{cls_name}) ensure_dir(out_dir) # 文件名 idx counters[cls_id] save_path os.path.join(out_dir, f{PREFIX}_{cls_name}_{idx:06d}.jpg) cv2.imwrite(save_path, crop) counters[cls_id] 1 print(处理完成) if __name__ __main__: main()1.6 制作osnet数据集在1.4生成的众多文件夹内同目录下创建以下脚本制作osnet数据集,更改其中的SRC_DIR,OUT_DIR# 所有种类是同种类时 import os import shutil import random # 配置 SRC_DIR classes3_output OUT_DIR reid_dataset_classes3 TRAIN_RATIO 0.7 # 70% 类别用于训练 random.seed(0) # def ensure_dir(path): if not os.path.exists(path): os.makedirs(path) def main(): classes [d for d in os.listdir(SRC_DIR) if os.path.isdir(os.path.join(SRC_DIR, d))] classes.sort() print(f总类别数: {len(classes)}) random.shuffle(classes) split_idx int(len(classes) * TRAIN_RATIO) train_classes classes[:split_idx] test_classes classes[split_idx:] # 创建输出目录 train_dir os.path.join(OUT_DIR, train) query_dir os.path.join(OUT_DIR, query) gallery_dir os.path.join(OUT_DIR, gallery) ensure_dir(train_dir) ensure_dir(query_dir) ensure_dir(gallery_dir) # 处理 train 类 train_id 1 for cls in train_classes: cls_path os.path.join(SRC_DIR, cls) imgs [f for f in os.listdir(cls_path) if f.endswith(.jpg)] for i, img_name in enumerate(imgs): cam_id 1 if i % 2 0 else 2 new_name f{train_id:04d}_c{cam_id:03d}_{i:08d}.jpg shutil.copy(os.path.join(cls_path, img_name), os.path.join(train_dir, new_name)) train_id 1 # 处理 query/gallery 类 test_id 1001 for cls in test_classes: cls_path os.path.join(SRC_DIR, cls) imgs [f for f in os.listdir(cls_path) if f.endswith(.jpg)] random.shuffle(imgs) # 至少两张用于 query/gallery query_img imgs[0] gallery_img imgs[1] shutil.copy(os.path.join(cls_path, query_img), os.path.join(query_dir, f{test_id:04d}_c001_00000000.jpg)) shutil.copy(os.path.join(cls_path, gallery_img), os.path.join(gallery_dir, f{test_id:04d}_c002_00000001.jpg)) # 剩余放 gallery for i, img_name in enumerate(imgs[2:]): new_name f{test_id:04d}_c002_{i2:08d}.jpg shutil.copy(os.path.join(cls_path, img_name), os.path.join(gallery_dir, new_name)) test_id 1 print(✅ 方案一数据集划分完成) if __name__ __main__: main()1.7 注册数据集osnet不支持像yolo那样以文件夹名为数据集需注册在deep-person-reid/torchreid/data/datasets/下新建如classes3.py文件内容如下更改其中的self.train_dir、self.query_dir、self.gallery_dir路径更改其中的dataset_dir更改其中的class classes3类名和对应的super(classes3, self).__init__(train, query, gallery, **kwargs)# classes3.py from __future__ import absolute_import import os.path as osp import glob import re from ..dataset import ImageDataset class classes3(ImageDataset): dataset_dir reid_dataset_classes3 # osnet的数据集路径 def __init__(self, root, **kwargs): self.root osp.abspath(osp.expanduser(root)) self.dataset_dir osp.join(self.root, self.dataset_dir) self.train_dir osp.join(self.dataset_dir, train) self.query_dir osp.join(self.dataset_dir, query) self.gallery_dir osp.join(self.dataset_dir, gallery) required_files [ self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir ] self.check_before_run(required_files) train self.process_dir(self.train_dir, relabelTrue) query self.process_dir(self.query_dir, relabelFalse) gallery self.process_dir(self.gallery_dir, relabelFalse) super(classes3, self).__init__(train, query, gallery, **kwargs) def process_dir(self, dir_path, relabelFalse): img_paths glob.glob(osp.join(dir_path, *.jpg)) pattern re.compile(r([-\d])_c(\d)) pid_container set() for img_path in img_paths: pid, _ map(int, pattern.search(img_path).groups()) if pid 0: continue pid_container.add(pid) pid2label {pid: label for label, pid in enumerate(pid_container)} data [] for img_path in img_paths: pid, camid map(int, pattern.search(img_path).groups()) if pid 0: continue camid - 1 # camid 从 0 开始 if relabel: pid pid2label[pid] data.append((img_path, pid, camid)) return data在 datasets/init.py 中注册文件路径model/deep-person-reid/torchreid/data/datasets/__init__.py添加:from .image.classes3 import classes3添加到 factory:找到__image_datasets {.....},在其中添加如下__image_datasets { ... classes3: classes3, }1.8 更改训练脚本train_osnet_3stage_auto.py完成后训练脚本的路径改为DATASET_NAME reid_dataset_classes3DATASET_ROOT改为classes3数据集的根路径2.yolo的冻结训练只需加一个参数freezefreeze 值含义22只训练 Detect Head极端小数据10Backbone 全冻结推荐起点15Backbone 部分 Neck0不冻结代码from ultralytics import YOLO import os os.environ[YOLO_DISABLE_AMP_CHECK] 1 model YOLO(./yolov8s.pt) # 加载模型 results model.train(datatruck.yaml, epochs50, workers16,imgsz640, batch64, device0,1,resumeFalse, freeze22, patience20)# cacheTrue # python -m torch.distributed.run --nproc_per_node2 ./train_detect.py
![[ERROR] Failed to convert PaddlePaddle model: (Unimplemented) the 0th elementwise MUST be ir::FloatA](http://pic.xiahunao.cn/yaotu/[ERROR] Failed to convert PaddlePaddle model: (Unimplemented) the 0th elementwise MUST be ir::FloatA)
