├── 28b589519f46fb619e634083b64003f4.png ├── 8aeca50f5a27ea09a2314486f8c82306.png ├── 9efd1dc903859a31c7a750ba5c460347.png ├── README.md ├── Thread_1.py ├── a9453f738b534a27cad4bee2f631a8dc.png ├── b5a81f82f8d6af35a8f284346b527a5f.png ├── c279a9fe26b6c3c735afe5b344dda0d1.png ├── d71134333984cc4b85e2072201ca308e.png ├── d9f36647fe3652388980467ef1939ac4.png └── train.py /28b589519f46fb619e634083b64003f4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/28b589519f46fb619e634083b64003f4.png -------------------------------------------------------------------------------- /8aeca50f5a27ea09a2314486f8c82306.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/8aeca50f5a27ea09a2314486f8c82306.png -------------------------------------------------------------------------------- /9efd1dc903859a31c7a750ba5c460347.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/9efd1dc903859a31c7a750ba5c460347.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 1.研究背景: 2 | 在传统城市道路交通信号灯的应用过程中，存在专业技术人员较缺乏、项目投入资金少、信号灯故障抢修较困难等问题，为城市居民生活与出行带来极大不便。对此，文章进行深入研究，分析了城市道路智慧交通信号灯系统的基本结构，并充分结合实际情况进行设计与应用，旨在缓解城市交通堵塞问题，提升居民幸福感。 3 | 4 | 5 | # 2.研究内容＆目标: 6 | 用python基于opencv实现对车辆流量的识别，然后通过yolov7算法计算南北和东西车道车数的比值，根据南北车道实时的车流量比值来对交通灯切换的时间比例进行实时指导，从而最大程度的保证城市交通的畅通，减少早晚高峰车辆拥堵的时间，实现智慧交通的设计目标。 7 | 8 | # 3.解决的关键问题: 9 | 1.理解交通灯切换的原理; 10 | 11 | 2.建立车辆检测、追踪和统计的模型; 12 | 13 | 3.根据假设建立实际问题优化模型。 14 | 15 | # 4.图片展示 16 | ![1.png](d71134333984cc4b85e2072201ca308e.png) 17 | 18 | ![2.png](28b589519f46fb619e634083b64003f4.png) 19 | 20 | 21 | # 5.视频展示 22 | [Python基于OpenCV的智能交通灯系统（南北车流量比例）[源码＆部署教程]](https://www.bilibili.com/video/BV1pd4y1X7AC?spm_id_from=333.999.0.0&vd_source=bc9aec86d164b67a7004b996143742dc) 23 | 24 | 25 | 26 | # 6.Deepsort目标追踪 27 | （1）获取原始视频帧 28 | （2）利用目标检测器对视频帧中的目标进行检测 29 | （3）将检测到的目标的框中的特征提取出来，该特征包括表观特征（方便特征对比避免ID switch）和运动特征（运动特征方 30 | 便卡尔曼滤波对其进行预测） 31 | （4）计算前后两帧目标之前的匹配程度（利用匈牙利算法和级联匹配），为每个追踪到的目标分配ID。 32 | Deepsort的前身是sort算法，sort算法的核心是卡尔曼滤波算法和匈牙利算法。 33 | 34 | 卡尔曼滤波算法作用：该算法的主要作用就是当前的一系列运动变量去预测下一时刻的运动变量，但是第一次的检测结果用来初始化卡尔曼滤波的运动变量。 35 | 36 | 匈牙利算法的作用：简单来讲就是解决分配问题，就是把一群检测框和卡尔曼预测的框做分配，让卡尔曼预测的框找到和自己最匹配的检测框，达到追踪的效果。 37 | 38 | #### sort工作流程如下图所示： 39 | ![4.png](a9453f738b534a27cad4bee2f631a8dc.png) 40 | 41 | Detections是通过目标检测到的框框。Tracks是轨迹信息。 42 | 43 | #### 整个算法的工作流程如下： 44 | （1）将第一帧检测到的结果创建其对应的Tracks。将卡尔曼滤波的运动变量初始化，通过卡尔曼滤波预测其对应的框框。 45 | 46 | （2）将该帧目标检测的框框和上一帧通过Tracks预测的框框一一进行IOU匹配，再通过IOU匹配的结果计算其代价矩阵（cost matrix，其计算方式是1-IOU）。 47 | 48 | （3）将（2）中得到的所有的代价矩阵作为匈牙利算法的输入，得到线性的匹配的结果，这时候我们得到的结果有三种，第一种是Tracks失配（Unmatched Tracks），我们直接将失配的Tracks删除；第二种是Detections失配（Unmatched Detections），我们将这样的Detections初始化为一个新的Tracks（new Tracks）；第三种是检测框和预测的框框配对成功，这说明我们前一帧和后一帧追踪成功，将其对应的Detections通过卡尔曼滤波更新其对应的Tracks变量。 49 | 50 | （4）反复循环（2）-（3）步骤，直到视频帧结束。 51 | 52 | #### Deepsort算法流程 53 | 由于sort算法还是比较粗糙的追踪算法，当物体发生遮挡的时候，特别容易丢失自己的ID。[该博客提出的方法在sort算法的基础上增加了级联匹配（Matching Cascade）和新轨迹的确认（confirmed）](https://afdian.net/item?plan_id=5db2f89862fd11eda8aa52540025c377)。Tracks分为确认态（confirmed），和不确认态（unconfirmed），新产生的Tracks是不确认态的；不确认态的Tracks必须要和Detections连续匹配一定的次数（默认是3）才可以转化成确认态。确认态的Tracks必须和Detections连续失配一定次数（默认30次），才会被删除。 54 | Deepsort算法的工作流程如下图所示： 55 | ![5.png](c279a9fe26b6c3c735afe5b344dda0d1.png) 56 | 整个算法的工作流程如下： 57 | 58 | （1）将第一帧次检测到的结果创建其对应的Tracks。将卡尔曼滤波的运动变量初始化，通过卡尔曼滤波预测其对应的框框。这时候的Tracks一定是unconfirmed的。 59 | 60 | （2）将该帧目标检测的框框和第上一帧通过Tracks预测的框框一一进行IOU匹配，再通过IOU匹配的结果计算其代价矩阵（cost matrix，其计算方式是1-IOU）。 61 | 62 | （3）将（2）中得到的所有的代价矩阵作为匈牙利算法的输入，得到线性的匹配的结果，这时候我们得到的结果有三种，第一种是Tracks失配（Unmatched Tracks），我们直接将失配的Tracks（因为这个Tracks是不确定态了，如果是确定态的话则要连续达到一定的次数（默认30次）才可以删除）删除；第二种是Detections失配（Unmatched Detections），我们将这样的Detections初始化为一个新的Tracks（new Tracks）；第三种是检测框和预测的框框配对成功，这说明我们前一帧和后一帧追踪成功，将其对应的Detections通过卡尔曼滤波更新其对应的Tracks变量。 63 | 64 | （4）反复循环（2）-（3）步骤，直到出现确认态（confirmed）的Tracks或者视频帧结束。 65 | 66 | （5）通过卡尔曼滤波预测其确认态的Tracks和不确认态的Tracks对应的框框。将确认态的Tracks的框框和是Detections进行级联匹配（之前每次只要Tracks匹配上都会保存Detections其的外观特征和运动信息，默认保存前100帧，利用外观特征和运动信息和Detections进行级联匹配,这么做是因为确认态（confirmed）的Tracks和Detections匹配的可能性更大）。 67 | 68 | （6）进行级联匹配后有三种可能的结果。第一种，Tracks匹配，这样的Tracks通过卡尔曼滤波更新其对应的Tracks变量。第二第三种是Detections和Tracks失配，这时将之前的不确认状态的Tracks和失配的Tracks一起和Unmatched Detections一一进行IOU匹配，再通过IOU匹配的结果计算其代价矩阵（cost matrix，其计算方式是1-IOU）。 69 | 70 | （7）将（6）中得到的所有的代价矩阵作为匈牙利算法的输入，得到线性的匹配的结果，这时候我们得到的结果有三种，第一种是Tracks失配（Unmatched Tracks），我们直接将失配的Tracks（因为这个Tracks是不确定态了，如果是确定态的话则要连续达到一定的次数（默认30次）才可以删除）删除；第二种是Detections失配（Unmatched Detections），我们将这样的Detections初始化为一个新的Tracks（new Tracks）；第三种是检测框和预测的框框配对成功，这说明我们前一帧和后一帧追踪成功，将其对应的Detections通过卡尔曼滤波更新其对应的Tracks变量。 71 | 72 | （8）反复循环（5）-（7）步骤，直到视频帧结束。 73 | 74 | 75 | 76 | # 7.准备YOLOv7格式数据集 77 | [如果不懂yolo格式数据集是什么样子的，建议先学习一下该博客](https://mbd.pub/o/bread/Y5Wbm5Zw)。大部分CVer都会推荐用labelImg进行数据的标注，我也不例外，推荐大家用labelImg进行数据标注。不过这里我不再详细介绍如何使用labelImg，网上有很多的教程。同时，标注数据需要用到图形交互界面，远程服务器就不太方便了，因此建议在本地电脑上标注好后再上传到服务器上。 78 | 79 | 这里假设我们已经得到标注好的yolo格式数据集，那么这个数据集将会按照如下的格式进行存放。 80 | ![n.png](b5a81f82f8d6af35a8f284346b527a5f.png) 81 | 不过在这里面，train_list.txt和val_list.txt是后来我们要自己生成的，而不是labelImg生成的；其他的则是labelImg生成的。 82 | 83 | 接下来，就是生成 train_list.txt和val_list.txt。train_list.txt存放了所有训练图片的路径，val_list.txt则是存放了所有验证图片的路径，如下图所示，一行代表一个图片的路径。这两个文件的生成写个循环就可以了，不算难。 84 | 85 | # 8.修改配置文件 86 | 总共有两个文件需要配置，一个是/yolov7/cfg/training/yolov7.yaml，这个文件是有关模型的配置文件；一个是/yolov7/data/coco.yaml，这个是数据集的配置文件。 87 | 88 | ## 第一步，复制yolov7.yaml文件到相同的路径下，然后重命名，我们重命名为yolov7-Helmet.yaml。 89 | 90 | ## 第二步，打开yolov7-Helmet.yaml文件，进行如下图所示的修改，这里修改的地方只有一处，就是把nc修改为我们数据集的目标总数即可。然后保存。 91 | 92 | ![b.png](d9f36647fe3652388980467ef1939ac4.png) 93 | 94 | ## 第三步，复制coco.yaml文件到相同的路径下，然后重命名，我们命名为Helmet.yaml。 95 | 96 | ## 第四步，打开Helmet.yaml文件，进行如下所示的修改，需要修改的地方为5处。 97 | 第一处：把代码自动下载COCO数据集的命令注释掉，以防代码自动下载数据集占用内存；第二处：修改train的位置为train_list.txt的路径；第三处：修改val的位置为val_list.txt的路径；第四处：修改nc为数据集目标总数；第五处：修改names为数据集所有目标的名称。然后保存。 98 | 99 | ![k.png](9efd1dc903859a31c7a750ba5c460347.png) 100 | 101 | # 9.训练代码 102 | ``` 103 | import argparse 104 | import logging 105 | import math 106 | import os 107 | import random 108 | import time 109 | from copy import deepcopy 110 | from pathlib import Path 111 | from threading import Thread 112 | 113 | import numpy as np 114 | import torch.distributed as dist 115 | import torch.nn as nn 116 | import torch.nn.functional as F 117 | import torch.optim as optim 118 | import torch.optim.lr_scheduler as lr_scheduler 119 | import torch.utils.data 120 | import yaml 121 | from torch.cuda import amp 122 | from torch.nn.parallel import DistributedDataParallel as DDP 123 | from torch.utils.tensorboard import SummaryWriter 124 | from tqdm import tqdm 125 | 126 | import test # import test.py to get mAP after each epoch 127 | from models.experimental import attempt_load 128 | from models.yolo import Model 129 | from utils.autoanchor import check_anchors 130 | from utils.datasets import create_dataloader 131 | from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \ 132 | fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \ 133 | check_requirements, print_mutation, set_logging, one_cycle, colorstr 134 | from utils.google_utils import attempt_download 135 | from utils.loss import ComputeLoss, ComputeLossOTA 136 | from utils.plots import plot_images, plot_labels, plot_results, plot_evolution 137 | from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel 138 | from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume 139 | 140 | logger = logging.getLogger(__name__) 141 | 142 | 143 | def train(hyp, opt, device, tb_writer=None): 144 | logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) 145 | save_dir, epochs, batch_size, total_batch_size, weights, rank, freeze = \ 146 | Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank, opt.freeze 147 | 148 | # Directories 149 | wdir = save_dir / 'weights' 150 | wdir.mkdir(parents=True, exist_ok=True) # make dir 151 | last = wdir / 'last.pt' 152 | best = wdir / 'best.pt' 153 | results_file = save_dir / 'results.txt' 154 | 155 | # Save run settings 156 | with open(save_dir / 'hyp.yaml', 'w') as f: 157 | yaml.dump(hyp, f, sort_keys=False) 158 | with open(save_dir / 'opt.yaml', 'w') as f: 159 | yaml.dump(vars(opt), f, sort_keys=False) 160 | 161 | # Configure 162 | plots = not opt.evolve # create plots 163 | cuda = device.type != 'cpu' 164 | init_seeds(2 + rank) 165 | with open(opt.data) as f: 166 | data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict 167 | is_coco = opt.data.endswith('coco.yaml') 168 | 169 | # Logging- Doing this before checking the dataset. Might update data_dict 170 | loggers = {'wandb': None} # loggers dict 171 | if rank in [-1, 0]: 172 | opt.hyp = hyp # add hyperparameters 173 | run_id = torch.load(weights, map_location=device).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None 174 | wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict) 175 | loggers['wandb'] = wandb_logger.wandb 176 | data_dict = wandb_logger.data_dict 177 | if wandb_logger.wandb: 178 | weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # WandbLogger might update weights, epochs if resuming 179 | 180 | nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes 181 | names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names 182 | assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check 183 | 184 | # Model 185 | pretrained = weights.endswith('.pt') 186 | if pretrained: 187 | with torch_distributed_zero_first(rank): 188 | attempt_download(weights) # download if not found locally 189 | ckpt = torch.load(weights, map_location=device) # load checkpoint 190 | model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 191 | exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys 192 | state_dict = ckpt['model'].float().state_dict() # to FP32 193 | state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect 194 | model.load_state_dict(state_dict, strict=False) # load 195 | logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report 196 | else: 197 | model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 198 | with torch_distributed_zero_first(rank): 199 | check_dataset(data_dict) # check 200 | train_path = data_dict['train'] 201 | test_path = data_dict['val'] 202 | 203 | # Freeze 204 | freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # parameter names to freeze (full or partial) 205 | for k, v in model.named_parameters(): 206 | v.requires_grad = True # train all layers 207 | if any(x in k for x in freeze): 208 | print('freezing %s' % k) 209 | v.requires_grad = False 210 | 211 | # Optimizer 212 | nbs = 64 # nominal batch size 213 | accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing 214 | hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay 215 | logger.info(f"Scaled weight_decay = {hyp['weight_decay']}") 216 | 217 | pg0, pg1, pg2 = [], [], [] # optimizer parameter groups 218 | for k, v in model.named_modules(): 219 | if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): 220 | pg2.append(v.bias) # biases 221 | if isinstance(v, nn.BatchNorm2d): 222 | pg0.append(v.weight) # no decay 223 | elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): 224 | pg1.append(v.weight) # apply decay 225 | if hasattr(v, 'im'): 226 | if hasattr(v.im, 'implicit'): 227 | pg0.append(v.im.implicit) 228 | else: 229 | for iv in v.im: 230 | pg0.append(iv.implicit) 231 | if hasattr(v, 'imc'): 232 | if hasattr(v.imc, 'implicit'): 233 | pg0.append(v.imc.implicit) 234 | else: 235 | for iv in v.imc: 236 | pg0.append(iv.implicit) 237 | if hasattr(v, 'imb'): 238 | if hasattr(v.imb, 'implicit'): 239 | pg0.append(v.imb.implicit) 240 | else: 241 | for iv in v.imb: 242 | pg0.append(iv.implicit) 243 | if hasattr(v, 'imo'): 244 | if hasattr(v.imo, 'implicit'): 245 | pg0.append(v.imo.implicit) 246 | else: 247 | for iv in v.imo: 248 | pg0.append(iv.implicit) 249 | if hasattr(v, 'ia'): 250 | if hasattr(v.ia, 'implicit'): 251 | pg0.append(v.ia.implicit) 252 | else: 253 | for iv in v.ia: 254 | pg0.append(iv.implicit) 255 | if hasattr(v, 'attn'): 256 | if hasattr(v.attn, 'logit_scale'): 257 | pg0.append(v.attn.logit_scale) 258 | if hasattr(v.attn, 'q_bias'): 259 | pg0.append(v.attn.q_bias) 260 | if hasattr(v.attn, 'v_bias'): 261 | pg0.append(v.attn.v_bias) 262 | if hasattr(v.attn, 'relative_position_bias_table'): 263 | pg0.append(v.attn.relative_position_bias_table) 264 | if hasattr(v, 'rbr_dense'): 265 | if hasattr(v.rbr_dense, 'weight_rbr_origin'): 266 | pg0.append(v.rbr_dense.weight_rbr_origin) 267 | if hasattr(v.rbr_dense, 'weight_rbr_avg_conv'): 268 | pg0.append(v.rbr_dense.weight_rbr_avg_conv) 269 | if hasattr(v.rbr_dense, 'weight_rbr_pfir_conv'): 270 | pg0.append(v.rbr_dense.weight_rbr_pfir_conv) 271 | if hasattr(v.rbr_dense, 'weight_rbr_1x1_kxk_idconv1'): 272 | pg0.append(v.rbr_dense.weight_rbr_1x1_kxk_idconv1) 273 | if hasattr(v.rbr_dense, 'weight_rbr_1x1_kxk_conv2'): 274 | pg0.append(v.rbr_dense.weight_rbr_1x1_kxk_conv2) 275 | if hasattr(v.rbr_dense, 'weight_rbr_gconv_dw'): 276 | pg0.append(v.rbr_dense.weight_rbr_gconv_dw) 277 | if hasattr(v.rbr_dense, 'weight_rbr_gconv_pw'): 278 | pg0.append(v.rbr_dense.weight_rbr_gconv_pw) 279 | if hasattr(v.rbr_dense, 'vector'): 280 | pg0.append(v.rbr_dense.vector) 281 | 282 | if opt.adam: 283 | optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum 284 | else: 285 | optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) 286 | 287 | optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay 288 | optimizer.add_param_group({'params': pg2}) # add pg2 (biases) 289 | logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0))) 290 | del pg0, pg1, pg2 291 | 292 | # Scheduler https://arxiv.org/pdf/1812.01187.pdf 293 | # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR 294 | if opt.linear_lr: 295 | lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear 296 | else: 297 | lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf'] 298 | scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) 299 | # plot_lr_scheduler(optimizer, scheduler, epochs) 300 | 301 | # EMA 302 | ema = ModelEMA(model) if rank in [-1, 0] else None 303 | 304 | # Resume 305 | start_epoch, best_fitness = 0, 0.0 306 | if pretrained: 307 | # Optimizer 308 | if ckpt['optimizer'] is not None: 309 | optimizer.load_state_dict(ckpt['optimizer']) 310 | best_fitness = ckpt['best_fitness'] 311 | 312 | # EMA 313 | if ema and ckpt.get('ema'): 314 | ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) 315 | ema.updates = ckpt['updates'] 316 | 317 | # Results 318 | if ckpt.get('training_results') is not None: 319 | results_file.write_text(ckpt['training_results']) # write results.txt 320 | 321 | # Epochs 322 | start_epoch = ckpt['epoch'] + 1 323 | if opt.resume: 324 | assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs) 325 | if epochs < start_epoch: 326 | logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' % 327 | (weights, ckpt['epoch'], epochs)) 328 | epochs += ckpt['epoch'] # finetune additional epochs 329 | 330 | del ckpt, state_dict 331 | 332 | # Image sizes 333 | gs = max(int(model.stride.max()), 32) # grid size (max stride) 334 | nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj']) 335 | imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples 336 | 337 | # DP mode 338 | if cuda and rank == -1 and torch.cuda.device_count() > 1: 339 | model = torch.nn.DataParallel(model) 340 | 341 | # SyncBatchNorm 342 | if opt.sync_bn and cuda and rank != -1: 343 | model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) 344 | logger.info('Using SyncBatchNorm()') 345 | 346 | # Trainloader 347 | dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, 348 | hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank, 349 | world_size=opt.world_size, workers=opt.workers, 350 | image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: ')) 351 | mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class 352 | nb = len(dataloader) # number of batches 353 | assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1) 354 | 355 | # Process 0 356 | if rank in [-1, 0]: 357 | testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader 358 | hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1, 359 | world_size=opt.world_size, workers=opt.workers, 360 | pad=0.5, prefix=colorstr('val: '))[0] 361 | 362 | if not opt.resume: 363 | labels = np.concatenate(dataset.labels, 0) 364 | c = torch.tensor(labels[:, 0]) # classes 365 | # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency 366 | # model._initialize_biases(cf.to(device)) 367 | if plots: 368 | #plot_labels(labels, names, save_dir, loggers) 369 | if tb_writer: 370 | tb_writer.add_histogram('classes', c, 0) 371 | 372 | # Anchors 373 | if not opt.noautoanchor: 374 | check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) 375 | model.half().float() # pre-reduce anchor precision 376 | 377 | # DDP mode 378 | if cuda and rank != -1: 379 | model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank, 380 | # nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698 381 | find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules())) 382 | 383 | # Model parameters 384 | hyp['box'] *= 3. / nl # scale to layers 385 | hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers 386 | hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers 387 | hyp['label_smoothing'] = opt.label_smoothing 388 | model.nc = nc # attach number of classes to model 389 | model.hyp = hyp # attach hyperparameters to model 390 | model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou) 391 | model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights 392 | model.names = names 393 | 394 | # Start training 395 | t0 = time.time() 396 | nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations) 397 | # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training 398 | maps = np.zeros(nc) # mAP per class 399 | results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) 400 | scheduler.last_epoch = start_epoch - 1 # do not move 401 | scaler = amp.GradScaler(enabled=cuda) 402 | compute_loss_ota = ComputeLossOTA(model) # init loss class 403 | compute_loss = ComputeLoss(model) # init loss class 404 | logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n' 405 | f'Using {dataloader.num_workers} dataloader workers\n' 406 | f'Logging results to {save_dir}\n' 407 | f'Starting training for {epochs} epochs...') 408 | torch.save(model, wdir / 'init.pt') 409 | for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ 410 | model.train() 411 | 412 | # Update image weights (optional) 413 | if opt.image_weights: 414 | # Generate indices 415 | if rank in [-1, 0]: 416 | cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights 417 | iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights 418 | dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx 419 | # Broadcast if DDP 420 | if rank != -1: 421 | indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int() 422 | dist.broadcast(indices, 0) 423 | if rank != 0: 424 | dataset.indices = indices.cpu().numpy() 425 | 426 | # Update mosaic border 427 | # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) 428 | # dataset.mosaic_border = [b - imgsz, -b] # height, width borders 429 | 430 | mloss = torch.zeros(4, device=device) # mean losses 431 | if rank != -1: 432 | dataloader.sampler.set_epoch(epoch) 433 | pbar = enumerate(dataloader) 434 | logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size')) 435 | if rank in [-1, 0]: 436 | pbar = tqdm(pbar, total=nb) # progress bar 437 | optimizer.zero_grad() 438 | for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- 439 | ni = i + nb * epoch # number integrated batches (since train start) 440 | imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0 441 | 442 | # Warmup 443 | if ni <= nw: 444 | xi = [0, nw] # x interp 445 | # model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou) 446 | accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round()) 447 | for j, x in enumerate(optimizer.param_groups): 448 | # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 449 | x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)]) 450 | if 'momentum' in x: 451 | x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) 452 | 453 | # Multi-scale 454 | if opt.multi_scale: 455 | sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size 456 | sf = sz / max(imgs.shape[2:]) # scale factor 457 | if sf != 1: 458 | ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) 459 | imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) 460 | 461 | # Forward 462 | with amp.autocast(enabled=cuda): 463 | pred = model(imgs) # forward 464 | if 'loss_ota' not in hyp or hyp['loss_ota'] == 1: 465 | loss, loss_items = compute_loss_ota(pred, targets.to(device), imgs) # loss scaled by batch_size 466 | else: 467 | loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size 468 | if rank != -1: 469 | loss *= opt.world_size # gradient averaged between devices in DDP mode 470 | if opt.quad: 471 | loss *= 4. 472 | 473 | # Backward 474 | scaler.scale(loss).backward() 475 | 476 | # Optimize 477 | if ni % accumulate == 0: 478 | scaler.step(optimizer) # optimizer.step 479 | scaler.update() 480 | optimizer.zero_grad() 481 | if ema: 482 | ema.update(model) 483 | 484 | # Print 485 | if rank in [-1, 0]: 486 | mloss = (mloss * i + loss_items) / (i + 1) # update mean losses 487 | mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB) 488 | s = ('%10s' * 2 + '%10.4g' * 6) % ( 489 | '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1]) 490 | pbar.set_description(s) 491 | 492 | # Plot 493 | if plots and ni < 10: 494 | f = save_dir / f'train_batch{ni}.jpg' # filename 495 | Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() 496 | # if tb_writer: 497 | # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) 498 | # tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph 499 | elif plots and ni == 10 and wandb_logger.wandb: 500 | wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in 501 | save_dir.glob('train*.jpg') if x.exists()]}) 502 | 503 | # end batch ------------------------------------------------------------------------------------------------ 504 | # end epoch ---------------------------------------------------------------------------------------------------- 505 | 506 | # Scheduler 507 | lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard 508 | scheduler.step() 509 | 510 | # DDP process 0 or single-GPU 511 | if rank in [-1, 0]: 512 | # mAP 513 | ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights']) 514 | final_epoch = epoch + 1 == epochs 515 | if not opt.notest or final_epoch: # Calculate mAP 516 | wandb_logger.current_epoch = epoch + 1 517 | results, maps, times = test.test(data_dict, 518 | batch_size=batch_size * 2, 519 | imgsz=imgsz_test, 520 | model=ema.ema, 521 | single_cls=opt.single_cls, 522 | dataloader=testloader, 523 | save_dir=save_dir, 524 | verbose=nc < 50 and final_epoch, 525 | plots=plots and final_epoch, 526 | wandb_logger=wandb_logger, 527 | compute_loss=compute_loss, 528 | is_coco=is_coco) 529 | 530 | # Write 531 | with open(results_file, 'a') as f: 532 | f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss 533 | if len(opt.name) and opt.bucket: 534 | os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name)) 535 | 536 | # Log 537 | tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss 538 | 'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 539 | 'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss 540 | 'x/lr0', 'x/lr1', 'x/lr2'] # params 541 | for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags): 542 | if tb_writer: 543 | tb_writer.add_scalar(tag, x, epoch) # tensorboard 544 | if wandb_logger.wandb: 545 | wandb_logger.log({tag: x}) # W&B 546 | 547 | # Update best mAP 548 | fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] 549 | if fi > best_fitness: 550 | best_fitness = fi 551 | wandb_logger.end_epoch(best_result=best_fitness == fi) 552 | 553 | # Save model 554 | if (not opt.nosave) or (final_epoch and not opt.evolve): # if save 555 | ckpt = {'epoch': epoch, 556 | 'best_fitness': best_fitness, 557 | 'training_results': results_file.read_text(), 558 | 'model': deepcopy(model.module if is_parallel(model) else model).half(), 559 | 'ema': deepcopy(ema.ema).half(), 560 | 'updates': ema.updates, 561 | 'optimizer': optimizer.state_dict(), 562 | 'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None} 563 | 564 | # Save last, best and delete 565 | torch.save(ckpt, last) 566 | if best_fitness == fi: 567 | torch.save(ckpt, best) 568 | if (best_fitness == fi) and (epoch >= 200): 569 | torch.save(ckpt, wdir / 'best_{:03d}.pt'.format(epoch)) 570 | if epoch == 0: 571 | torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) 572 | elif ((epoch+1) % 25) == 0: 573 | torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) 574 | elif epoch >= (epochs-5): 575 | torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) 576 | if wandb_logger.wandb: 577 | if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1: 578 | wandb_logger.log_model( 579 | last.parent, opt, epoch, fi, best_model=best_fitness == fi) 580 | del ckpt 581 | 582 | # end epoch ---------------------------------------------------------------------------------------------------- 583 | # end training 584 | if rank in [-1, 0]: 585 | # Plots 586 | if plots: 587 | plot_results(save_dir=save_dir) # save as results.png 588 | if wandb_logger.wandb: 589 | files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]] 590 | wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files 591 | if (save_dir / f).exists()]}) 592 | # Test best.pt 593 | logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600)) 594 | if opt.data.endswith('coco.yaml') and nc == 80: # if COCO 595 | for m in (last, best) if best.exists() else (last): # speed, mAP tests 596 | results, _, _ = test.test(opt.data, 597 | batch_size=batch_size * 2, 598 | imgsz=imgsz_test, 599 | conf_thres=0.001, 600 | iou_thres=0.7, 601 | model=attempt_load(m, device).half(), 602 | single_cls=opt.single_cls, 603 | dataloader=testloader, 604 | save_dir=save_dir, 605 | save_json=True, 606 | plots=False, 607 | is_coco=is_coco) 608 | 609 | # Strip optimizers 610 | final = best if best.exists() else last # final model 611 | for f in last, best: 612 | if f.exists(): 613 | strip_optimizer(f) # strip optimizers 614 | if opt.bucket: 615 | os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload 616 | if wandb_logger.wandb and not opt.evolve: # Log the stripped model 617 | wandb_logger.wandb.log_artifact(str(final), type='model', 618 | name='run_' + wandb_logger.wandb_run.id + '_model', 619 | aliases=['last', 'best', 'stripped']) 620 | wandb_logger.finish_run() 621 | else: 622 | dist.destroy_process_group() 623 | torch.cuda.empty_cache() 624 | return results 625 | 626 | 627 | if __name__ == '__main__': 628 | parser = argparse.ArgumentParser() 629 | parser.add_argument('--weights', type=str, default='yolov7.pt', help='initial weights path') 630 | parser.add_argument('--cfg', type=str, default='cfg/training/yolov7.yaml', help='model.yaml path') 631 | parser.add_argument('--data', type=str, default='data/coco.yaml', help='data.yaml path') 632 | parser.add_argument('--hyp', type=str, default='data/hyp.scratch.p5.yaml', help='hyperparameters path') 633 | parser.add_argument('--epochs', type=int, default=300) 634 | parser.add_argument('--batch-size', type=int, default=4, help='total batch size for all GPUs') 635 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') 636 | parser.add_argument('--rect', action='store_true', help='rectangular training') 637 | parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') 638 | parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') 639 | parser.add_argument('--notest', action='store_true', help='only test final epoch') 640 | parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check') 641 | parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') 642 | parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') 643 | parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') 644 | parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') 645 | parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 646 | parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') 647 | parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class') 648 | parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') 649 | parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') 650 | parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') 651 | parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers') 652 | parser.add_argument('--project', default='runs/train', help='save to project/name') 653 | parser.add_argument('--entity', default=None, help='W&B entity') 654 | parser.add_argument('--name', default='exp', help='save to project/name') 655 | parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') 656 | parser.add_argument('--quad', action='store_true', help='quad dataloader') 657 | parser.add_argument('--linear-lr', action='store_true', help='linear LR') 658 | parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon') 659 | parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table') 660 | parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B') 661 | parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch') 662 | parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used') 663 | parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone of yolov7=50, first3=0 1 2') 664 | opt = parser.parse_args() 665 | 666 | # Set DDP variables 667 | opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1 668 | opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1 669 | set_logging(opt.global_rank) 670 | #if opt.global_rank in [-1, 0]: 671 | # check_git_status() 672 | # check_requirements() 673 | 674 | # Resume 675 | wandb_run = check_wandb_resume(opt) 676 | if opt.resume and not wandb_run: # resume an interrupted run 677 | ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path 678 | assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' 679 | apriori = opt.global_rank, opt.local_rank 680 | with open(Path(ckpt).parent.parent / 'opt.yaml') as f: 681 | opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader)) # replace 682 | opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori # reinstate 683 | logger.info('Resuming training from %s' % ckpt) 684 | else: 685 | # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml') 686 | opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp) # check files 687 | assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified' 688 | opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test) 689 | opt.name = 'evolve' if opt.evolve else opt.name 690 | opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve) # increment run 691 | 692 | # DDP mode 693 | opt.total_batch_size = opt.batch_size 694 | device = select_device(opt.device, batch_size=opt.batch_size) 695 | if opt.local_rank != -1: 696 | assert torch.cuda.device_count() > opt.local_rank 697 | torch.cuda.set_device(opt.local_rank) 698 | device = torch.device('cuda', opt.local_rank) 699 | dist.init_process_group(backend='nccl', init_method='env://') # distributed backend 700 | assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count' 701 | opt.batch_size = opt.total_batch_size // opt.world_size 702 | 703 | # Hyperparameters 704 | with open(opt.hyp) as f: 705 | hyp = yaml.load(f, Loader=yaml.SafeLoader) # load hyps 706 | 707 | # Train 708 | logger.info(opt) 709 | if not opt.evolve: 710 | tb_writer = None # init loggers 711 | if opt.global_rank in [-1, 0]: 712 | prefix = colorstr('tensorboard: ') 713 | logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/") 714 | tb_writer = SummaryWriter(opt.save_dir) # Tensorboard 715 | train(hyp, opt, device, tb_writer) 716 | 717 | # Evolve hyperparameters (optional) 718 | else: 719 | # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) 720 | meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) 721 | 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) 722 | 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 723 | 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay 724 | 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok) 725 | 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum 726 | 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr 727 | 'box': (1, 0.02, 0.2), # box loss gain 728 | 'cls': (1, 0.2, 4.0), # cls loss gain 729 | 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight 730 | 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) 731 | 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight 732 | 'iou_t': (0, 0.1, 0.7), # IoU training threshold 733 | 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold 734 | 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore) 735 | 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) 736 | 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) 737 | 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) 738 | 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) 739 | 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) 740 | 'translate': (1, 0.0, 0.9), # image translation (+/- fraction) 741 | 'scale': (1, 0.0, 0.9), # image scale (+/- gain) 742 | 'shear': (1, 0.0, 10.0), # image shear (+/- deg) 743 | 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 744 | 'flipud': (1, 0.0, 1.0), # image flip up-down (probability) 745 | 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) 746 | 'mosaic': (1, 0.0, 1.0), # image mixup (probability) 747 | 'mixup': (1, 0.0, 1.0), # image mixup (probability) 748 | 'copy_paste': (1, 0.0, 1.0), # segment copy-paste (probability) 749 | 'paste_in': (1, 0.0, 1.0)} # segment copy-paste (probability) 750 | 751 | with open(opt.hyp, errors='ignore') as f: 752 | hyp = yaml.safe_load(f) # load hyps dict 753 | if 'anchors' not in hyp: # anchors commented in hyp.yaml 754 | hyp['anchors'] = 3 755 | 756 | assert opt.local_rank == -1, 'DDP mode not implemented for --evolve' 757 | opt.notest, opt.nosave = True, True # only test/save final epoch 758 | # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices 759 | yaml_file = Path(opt.save_dir) / 'hyp_evolved.yaml' # save best result here 760 | if opt.bucket: 761 | os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket) # download evolve.txt if exists 762 | 763 | for _ in range(300): # generations to evolve 764 | if Path('evolve.txt').exists(): # if evolve.txt exists: select best hyps and mutate 765 | # Select parent(s) 766 | parent = 'single' # parent selection method: 'single' or 'weighted' 767 | x = np.loadtxt('evolve.txt', ndmin=2) 768 | n = min(5, len(x)) # number of previous results to consider 769 | x = x[np.argsort(-fitness(x))][:n] # top n mutations 770 | w = fitness(x) - fitness(x).min() # weights 771 | if parent == 'single' or len(x) == 1: 772 | # x = x[random.randint(0, n - 1)] # random selection 773 | x = x[random.choices(range(n), weights=w)[0]] # weighted selection 774 | elif parent == 'weighted': 775 | x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination 776 | 777 | # Mutate 778 | mp, s = 0.8, 0.2 # mutation probability, sigma 779 | npr = np.random 780 | npr.seed(int(time.time())) 781 | g = np.array([x[0] for x in meta.values()]) # gains 0-1 782 | ng = len(meta) 783 | v = np.ones(ng) 784 | while all(v == 1): # mutate until a change occurs (prevent duplicates) 785 | v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0) 786 | for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300) 787 | hyp[k] = float(x[i + 7] * v[i]) # mutate 788 | 789 | # Constrain to limits 790 | for k, v in meta.items(): 791 | hyp[k] = max(hyp[k], v[1]) # lower limit 792 | hyp[k] = min(hyp[k], v[2]) # upper limit 793 | hyp[k] = round(hyp[k], 5) # significant digits 794 | 795 | # Train mutation 796 | results = train(hyp.copy(), opt, device) 797 | 798 | # Write mutation results 799 | print_mutation(hyp.copy(), results, yaml_file, opt.bucket) 800 | 801 | # Plot results 802 | plot_evolution(yaml_file) 803 | print(f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n' 804 | f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}') 805 | ``` 806 | # 10.UI界面的编写＆系统的整合 807 | ``` 808 | class Thread_1(QThread): # 线程1 809 | def __init__(self,info1): 810 | super().__init__() 811 | self.info1=info1 812 | self.run2(self.info1) 813 | 814 | def run2(self, info1): 815 | result = [] 816 | result = det_yolov7(info1) 817 | 818 | 819 | class Ui_MainWindow(object): 820 | def setupUi(self, MainWindow): 821 | MainWindow.setObjectName("MainWindow") 822 | MainWindow.resize(1280, 960) 823 | MainWindow.setStyleSheet("background-image: url(\"./template/carui.png\")") 824 | self.centralwidget = QtWidgets.QWidget(MainWindow) 825 | self.centralwidget.setObjectName("centralwidget") 826 | self.label = QtWidgets.QLabel(self.centralwidget) 827 | self.label.setGeometry(QtCore.QRect(168, 60, 551, 71)) 828 | self.label.setAutoFillBackground(False) 829 | self.label.setStyleSheet("") 830 | self.label.setFrameShadow(QtWidgets.QFrame.Plain) 831 | self.label.setAlignment(QtCore.Qt.AlignCenter) 832 | self.label.setObjectName("label") 833 | self.label.setStyleSheet("font-size:42px;font-weight:bold;font-family:SimHei;background:rgba(255,255,255,0);") 834 | self.label_2 = QtWidgets.QLabel(self.centralwidget) 835 | self.label_2.setGeometry(QtCore.QRect(40, 188, 751, 501)) 836 | self.label_2.setStyleSheet("background:rgba(255,255,255,1);") 837 | self.label_2.setAlignment(QtCore.Qt.AlignCenter) 838 | self.label_2.setObjectName("label_2") 839 | self.textBrowser = QtWidgets.QTextBrowser(self.centralwidget) 840 | self.textBrowser.setGeometry(QtCore.QRect(73, 746, 851, 174)) 841 | self.textBrowser.setStyleSheet("background:rgba(0,0,0,0);") 842 | self.textBrowser.setObjectName("textBrowser") 843 | self.pushButton = QtWidgets.QPushButton(self.centralwidget) 844 | self.pushButton.setGeometry(QtCore.QRect(1020, 750, 150, 40)) 845 | self.pushButton.setStyleSheet("background:rgba(53,142,255,1);border-radius:10px;padding:2px 4px;") 846 | self.pushButton.setObjectName("pushButton") 847 | self.pushButton_2 = QtWidgets.QPushButton(self.centralwidget) 848 | self.pushButton_2.setGeometry(QtCore.QRect(1020, 810, 150, 40)) 849 | self.pushButton_2.setStyleSheet("background:rgba(53,142,255,1);border-radius:10px;padding:2px 4px;") 850 | self.pushButton_2.setObjectName("pushButton_2") 851 | self.pushButton_3 = QtWidgets.QPushButton(self.centralwidget) 852 | self.pushButton_3.setGeometry(QtCore.QRect(1020, 870, 150, 40)) 853 | self.pushButton_3.setStyleSheet("background:rgba(53,142,255,1);border-radius:10px;padding:2px 4px;") 854 | self.pushButton_3.setObjectName("pushButton_2") 855 | MainWindow.setCentralWidget(self.centralwidget) 856 | 857 | self.retranslateUi(MainWindow) 858 | QtCore.QMetaObject.connectSlotsByName(MainWindow) 859 | 860 | def retranslateUi(self, MainWindow): 861 | _translate = QtCore.QCoreApplication.translate 862 | MainWindow.setWindowTitle(_translate("MainWindow", "基于YOLO＆Deepsort的交通车流量统计系统")) 863 | self.label.setText(_translate("MainWindow", "基于YOLO＆Deepsort的交通车流量统计系统")) 864 | self.label_2.setText(_translate("MainWindow", "请添加对象，注意路径不要存在中文")) 865 | self.pushButton.setText(_translate("MainWindow", "选择对象")) 866 | self.pushButton_2.setText(_translate("MainWindow", "开始识别")) 867 | self.pushButton_3.setText(_translate("MainWindow", "退出系统")) 868 | 869 | # 点击文本框绑定槽事件 870 | self.pushButton.clicked.connect(self.openfile) 871 | self.pushButton_2.clicked.connect(self.click_1) 872 | self.pushButton_3.clicked.connect(self.handleCalc3) 873 | 874 | def openfile(self): 875 | global sname, filepath 876 | fname = QFileDialog() 877 | fname.setAcceptMode(QFileDialog.AcceptOpen) 878 | fname, _ = fname.getOpenFileName() 879 | if fname == '': 880 | return 881 | filepath = os.path.normpath(fname) 882 | sname = filepath.split(os.sep) 883 | ui.printf("当前选择的文件路径是：%s" % filepath) 884 | try: 885 | show = cv2.imread(filepath) 886 | ui.showimg(show) 887 | except: 888 | ui.printf('请检查路径是否存在中文，更名后重试！') 889 | 890 | 891 | def handleCalc3(self): 892 | os._exit(0) 893 | 894 | def printf(self,text): 895 | self.textBrowser.append(text) 896 | self.cursor = self.textBrowser.textCursor() 897 | self.textBrowser.moveCursor(self.cursor.End) 898 | QtWidgets.QApplication.processEvents() 899 | 900 | def showimg(self,img): 901 | global vid 902 | img2 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 903 | 904 | _image = QtGui.QImage(img2[:], img2.shape[1], img2.shape[0], img2.shape[1] * 3, 905 | QtGui.QImage.Format_RGB888) 906 | n_width = _image.width() 907 | n_height = _image.height() 908 | if n_width / 500 >= n_height / 400: 909 | ratio = n_width / 700 910 | else: 911 | ratio = n_height / 700 912 | new_width = int(n_width / ratio) 913 | new_height = int(n_height / ratio) 914 | new_img = _image.scaled(new_width, new_height, Qt.KeepAspectRatio) 915 | self.label_2.setPixmap(QPixmap.fromImage(new_img)) 916 | 917 | def click_1(self): 918 | global filepath 919 | try: 920 | self.thread_1.quit() 921 | except: 922 | pass 923 | self.thread_1 = Thread_1(filepath) # 创建线程 924 | self.thread_1.wait() 925 | self.thread_1.start() # 开始线程 926 | 927 | 928 | if __name__ == "__main__": 929 | app = QtWidgets.QApplication(sys.argv) 930 | MainWindow = QtWidgets.QMainWindow() 931 | ui = Ui_MainWindow() 932 | ui.setupUi(MainWindow) 933 | MainWindow.show() 934 | sys.exit(app.exec_()) 935 | ``` 936 | 937 | 938 | # 11.项目展示 939 | 下图[完整源码&环境部署视频教程&自定义UI界面＆技术文档](https://s.xiaocichang.com/s/db011f) 940 | ![4.png](8aeca50f5a27ea09a2314486f8c82306.png) 941 | 参考博客[《Python基于OpenCV的智能交通灯系统（南北车流量比例）\[源码＆部署教程\]》](https://mbd.pub/o/qunma/work) 942 | 943 | 944 | 945 | 946 | 947 | 948 | 949 | # 12.参考文献: 950 | *** 951 | * *[1]*[基于北斗三号导航系统的卫星定位技术](https://kns.cnki.net/kcms/detail/detail.aspx?filename=ZGKJ202202008&dbcode=CJFD&dbname=CJFDTEMP&v=GomzuAQ4SIVMgl7fZ7_yDabHpEJtfvAChsLv-Va1wlKSHhxc1_2eHkLBYtgQqNGM)[J]. 赵鹏飞,陈高峰,李小娟,汪大宝,刘宁,王慧聪. 中国空间科学技术. 2022(02) 952 | * *[2]*[低轨卫星导航增强系统精密星历设计方法研究](https://kns.cnki.net/kcms/detail/detail.aspx?filename=DKXB202108006&dbcode=CJFD&dbname=CJFDTEMP&v=ZEGokJOHMJornwXCzXauXtKUhfggIUfOgM7X7C84_yndDVlk2QlOds_hWOpa7N8t)[J]. 郭学立,王磊. 大地测量与地球动力学. 2021(08) 953 | * *[3]*[一种改进的深度确定性策略梯度网络交通信号控制系统](https://kns.cnki.net/kcms/detail/detail.aspx?filename=SCDX202104012&dbcode=CJFD&dbname=CJFDTEMP&v=0B8uIgdMGrZJOpTVo9AEGZavv8XixJKNw3cUaqWfQTYloRnThagxSOgiICqUJ4k3)[J]. 刘利军,王州,余臻. 四川大学学报(自然科学版). 2021(04) 954 | * *[4]*[城市道路照明工程电气设计研究](https://kns.cnki.net/kcms/detail/detail.aspx?filename=GYZM202106002&dbcode=CJFD&dbname=CJFDTEMP&v=eBc_zmd0FBFpvq6Ceu4FyCBs4B0ZtdKHlyyuFi2hi_0nYkV4FniB8-pYvJHvGdIY)[J]. 曾海宁. 光源与照明. 2021(06) 955 | * *[5]*[超高精度定位](https://kns.cnki.net/kcms/detail/detail.aspx?filename=ZKJJ202103014&dbcode=CJFD&dbname=CJFD2021&v=vVIQLcBp36cpzTGqWnp_2SLNUZqTJk8O2_ye3xrqNQs5i4IHEStdig_ESsHsMyry)[J]. 杨元喜,任夏. 中国科学基金. 2021(03) 956 | * *[6]*[主干道信号灯通行可靠度量化评估仿真](https://kns.cnki.net/kcms/detail/detail.aspx?filename=JSJZ202103015&dbcode=CJFD&dbname=CJFDTEMP&v=hcKC0M6OGdJoDZqVc0FPIxP54rVDxV3CAg-Xf5NbF5pvgF0qRGayREWwrK1tMteD)[J]. 薛煜亮,郝艳召. 计算机仿真. 2021(03) 957 | * *[7]*[基于自学习的组网式交通信号灯异常检测研究](https://kns.cnki.net/kcms/detail/detail.aspx?filename=CQJT202103005&dbcode=CJFD&dbname=CJFDTEMP&v=gHNteUSxkIZmnsVTwR3A-Vo_bVSMahSwtuYSJoa4_FFCvWTbKN2RyFKZ-31icBha)[J]. 刘永涛,樊亚敏,张莉,黎冠. 重庆交通大学学报(自然科学版). 2021(03) 958 | * *[8]*[采用物联网技术的智慧交通灯控制系统设计](https://kns.cnki.net/kcms/detail/detail.aspx?filename=LMZY202004018&dbcode=CJFD&dbname=CJFDTEMP&v=T7DkGeD04UE5pNwGzkTaHGEFROT0y2HqMUnMJzENMmmvSlH3DQEuIYPoi4v3WiOJ)[J]. 施凤鸣,骆旭坤. 黎明职业大学学报. 2020(04) 959 | * *[9]*[交叉口交通信号灯的模糊控制及优化研究](https://kns.cnki.net/kcms/detail/detail.aspx?filename=XTFZ202012013&dbcode=DKFX&dbname=DKFXTEMP&v=eLe5N9-vhWJWT57q2hMh_LKsL6BZ8ZcUmfy_eE528jnfatJMyYLOrkhOAuvcTiS5)[J]. 刘佳佳,左兴权. 系统仿真学报. 2020(12) 960 | * *[10]*[交通信号灯线路漏电及保护技术分析](https://kns.cnki.net/kcms/detail/detail.aspx?filename=YJCO202014054&dbcode=CJFD&dbname=CJFDTEMP&v=DRf6h3yqiFkAaAkGOj-pb-S0y03PM8r5GkPXaauCapX7noNniui3XdZTXAnKapAa)[J]. 冀光普,余仁辉,冯国行. 工程技术研究. 2020(14) 961 | 962 | 963 | 964 | 965 | 966 | 967 | 968 | 969 | --- 970 | #### 如果您需要更详细的【源码和环境部署教程】，除了通过【系统整合】小节的链接获取之外，还可以通过邮箱以下途径获取: 971 | #### 1.请先在GitHub上为该项目点赞（Star），编辑一封邮件，附上点赞的截图、项目的中文描述概述（About）以及您的用途需求，发送到我们的邮箱 972 | #### sharecode@yeah.net 973 | #### 2.我们收到邮件后会定期根据邮件的接收顺序将【完整源码和环境部署教程】发送到您的邮箱。 974 | #### 【免责声明】本文来源于用户投稿，如果侵犯任何第三方的合法权益，可通过邮箱联系删除。 -------------------------------------------------------------------------------- /Thread_1.py: -------------------------------------------------------------------------------- 1 | 2 | class Thread_1(QThread): # 线程1 3 | def __init__(self,info1): 4 | super().__init__() 5 | self.info1=info1 6 | self.run2(self.info1) 7 | 8 | def run2(self, info1): 9 | result = [] 10 | result = det_yolov7(info1) 11 | 12 | 13 | class Ui_MainWindow(object): 14 | def setupUi(self, MainWindow): 15 | MainWindow.setObjectName("MainWindow") 16 | MainWindow.resize(1280, 960) 17 | MainWindow.setStyleSheet("background-image: url(\"./template/carui.png\")") 18 | self.centralwidget = QtWidgets.QWidget(MainWindow) 19 | self.centralwidget.setObjectName("centralwidget") 20 | self.label = QtWidgets.QLabel(self.centralwidget) 21 | self.label.setGeometry(QtCore.QRect(168, 60, 551, 71)) 22 | self.label.setAutoFillBackground(False) 23 | self.label.setStyleSheet("") 24 | self.label.setFrameShadow(QtWidgets.QFrame.Plain) 25 | self.label.setAlignment(QtCore.Qt.AlignCenter) 26 | self.label.setObjectName("label") 27 | self.label.setStyleSheet("font-size:42px;font-weight:bold;font-family:SimHei;background:rgba(255,255,255,0);") 28 | self.label_2 = QtWidgets.QLabel(self.centralwidget) 29 | self.label_2.setGeometry(QtCore.QRect(40, 188, 751, 501)) 30 | self.label_2.setStyleSheet("background:rgba(255,255,255,1);") 31 | self.label_2.setAlignment(QtCore.Qt.AlignCenter) 32 | self.label_2.setObjectName("label_2") 33 | self.textBrowser = QtWidgets.QTextBrowser(self.centralwidget) 34 | self.textBrowser.setGeometry(QtCore.QRect(73, 746, 851, 174)) 35 | self.textBrowser.setStyleSheet("background:rgba(0,0,0,0);") 36 | self.textBrowser.setObjectName("textBrowser") 37 | self.pushButton = QtWidgets.QPushButton(self.centralwidget) 38 | self.pushButton.setGeometry(QtCore.QRect(1020, 750, 150, 40)) 39 | self.pushButton.setStyleSheet("background:rgba(53,142,255,1);border-radius:10px;padding:2px 4px;") 40 | self.pushButton.setObjectName("pushButton") 41 | self.pushButton_2 = QtWidgets.QPushButton(self.centralwidget) 42 | self.pushButton_2.setGeometry(QtCore.QRect(1020, 810, 150, 40)) 43 | self.pushButton_2.setStyleSheet("background:rgba(53,142,255,1);border-radius:10px;padding:2px 4px;") 44 | self.pushButton_2.setObjectName("pushButton_2") 45 | self.pushButton_3 = QtWidgets.QPushButton(self.centralwidget) 46 | self.pushButton_3.setGeometry(QtCore.QRect(1020, 870, 150, 40)) 47 | self.pushButton_3.setStyleSheet("background:rgba(53,142,255,1);border-radius:10px;padding:2px 4px;") 48 | self.pushButton_3.setObjectName("pushButton_2") 49 | MainWindow.setCentralWidget(self.centralwidget) 50 | 51 | self.retranslateUi(MainWindow) 52 | QtCore.QMetaObject.connectSlotsByName(MainWindow) 53 | 54 | def retranslateUi(self, MainWindow): 55 | _translate = QtCore.QCoreApplication.translate 56 | MainWindow.setWindowTitle(_translate("MainWindow", "基于YOLO＆Deepsort的交通车流量统计系统")) 57 | self.label.setText(_translate("MainWindow", "基于YOLO＆Deepsort的交通车流量统计系统")) 58 | self.label_2.setText(_translate("MainWindow", "请添加对象，注意路径不要存在中文")) 59 | self.pushButton.setText(_translate("MainWindow", "选择对象")) 60 | self.pushButton_2.setText(_translate("MainWindow", "开始识别")) 61 | self.pushButton_3.setText(_translate("MainWindow", "退出系统")) 62 | 63 | # 点击文本框绑定槽事件 64 | self.pushButton.clicked.connect(self.openfile) 65 | self.pushButton_2.clicked.connect(self.click_1) 66 | self.pushButton_3.clicked.connect(self.handleCalc3) 67 | 68 | def openfile(self): 69 | global sname, filepath 70 | fname = QFileDialog() 71 | fname.setAcceptMode(QFileDialog.AcceptOpen) 72 | fname, _ = fname.getOpenFileName() 73 | if fname == '': 74 | return 75 | filepath = os.path.normpath(fname) 76 | sname = filepath.split(os.sep) 77 | ui.printf("当前选择的文件路径是：%s" % filepath) 78 | try: 79 | show = cv2.imread(filepath) 80 | ui.showimg(show) 81 | except: 82 | ui.printf('请检查路径是否存在中文，更名后重试！') 83 | 84 | 85 | def handleCalc3(self): 86 | os._exit(0) 87 | 88 | def printf(self,text): 89 | self.textBrowser.append(text) 90 | self.cursor = self.textBrowser.textCursor() 91 | self.textBrowser.moveCursor(self.cursor.End) 92 | QtWidgets.QApplication.processEvents() 93 | 94 | def showimg(self,img): 95 | global vid 96 | img2 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 97 | 98 | _image = QtGui.QImage(img2[:], img2.shape[1], img2.shape[0], img2.shape[1] * 3, 99 | QtGui.QImage.Format_RGB888) 100 | n_width = _image.width() 101 | n_height = _image.height() 102 | if n_width / 500 >= n_height / 400: 103 | ratio = n_width / 700 104 | else: 105 | ratio = n_height / 700 106 | new_width = int(n_width / ratio) 107 | new_height = int(n_height / ratio) 108 | new_img = _image.scaled(new_width, new_height, Qt.KeepAspectRatio) 109 | self.label_2.setPixmap(QPixmap.fromImage(new_img)) 110 | 111 | def click_1(self): 112 | global filepath 113 | try: 114 | self.thread_1.quit() 115 | except: 116 | pass 117 | self.thread_1 = Thread_1(filepath) # 创建线程 118 | self.thread_1.wait() 119 | self.thread_1.start() # 开始线程 120 | 121 | 122 | if __name__ == "__main__": 123 | app = QtWidgets.QApplication(sys.argv) 124 | MainWindow = QtWidgets.QMainWindow() 125 | ui = Ui_MainWindow() 126 | ui.setupUi(MainWindow) 127 | MainWindow.show() 128 | sys.exit(app.exec_()) 129 | -------------------------------------------------------------------------------- /a9453f738b534a27cad4bee2f631a8dc.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/a9453f738b534a27cad4bee2f631a8dc.png -------------------------------------------------------------------------------- /b5a81f82f8d6af35a8f284346b527a5f.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/b5a81f82f8d6af35a8f284346b527a5f.png -------------------------------------------------------------------------------- /c279a9fe26b6c3c735afe5b344dda0d1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/c279a9fe26b6c3c735afe5b344dda0d1.png -------------------------------------------------------------------------------- /d71134333984cc4b85e2072201ca308e.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/d71134333984cc4b85e2072201ca308e.png -------------------------------------------------------------------------------- /d9f36647fe3652388980467ef1939ac4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qunshansj/intelligent-traffic-light-system-opencv-python/0af96efd31afa163a03c133fd7ed111aafceb824/d9f36647fe3652388980467ef1939ac4.png -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | 2 | import argparse 3 | import logging 4 | import math 5 | import os 6 | import random 7 | import time 8 | from copy import deepcopy 9 | from pathlib import Path 10 | from threading import Thread 11 | 12 | import numpy as np 13 | import torch.distributed as dist 14 | import torch.nn as nn 15 | import torch.nn.functional as F 16 | import torch.optim as optim 17 | import torch.optim.lr_scheduler as lr_scheduler 18 | import torch.utils.data 19 | import yaml 20 | from torch.cuda import amp 21 | from torch.nn.parallel import DistributedDataParallel as DDP 22 | from torch.utils.tensorboard import SummaryWriter 23 | from tqdm import tqdm 24 | 25 | import test # import test.py to get mAP after each epoch 26 | from models.experimental import attempt_load 27 | from models.yolo import Model 28 | from utils.autoanchor import check_anchors 29 | from utils.datasets import create_dataloader 30 | from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \ 31 | fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \ 32 | check_requirements, print_mutation, set_logging, one_cycle, colorstr 33 | from utils.google_utils import attempt_download 34 | from utils.loss import ComputeLoss, ComputeLossOTA 35 | from utils.plots import plot_images, plot_labels, plot_results, plot_evolution 36 | from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel 37 | from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume 38 | 39 | logger = logging.getLogger(__name__) 40 | 41 | 42 | def train(hyp, opt, device, tb_writer=None): 43 | logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) 44 | save_dir, epochs, batch_size, total_batch_size, weights, rank, freeze = \ 45 | Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank, opt.freeze 46 | 47 | # Directories 48 | wdir = save_dir / 'weights' 49 | wdir.mkdir(parents=True, exist_ok=True) # make dir 50 | last = wdir / 'last.pt' 51 | best = wdir / 'best.pt' 52 | results_file = save_dir / 'results.txt' 53 | 54 | # Save run settings 55 | with open(save_dir / 'hyp.yaml', 'w') as f: 56 | yaml.dump(hyp, f, sort_keys=False) 57 | with open(save_dir / 'opt.yaml', 'w') as f: 58 | yaml.dump(vars(opt), f, sort_keys=False) 59 | 60 | # Configure 61 | plots = not opt.evolve # create plots 62 | cuda = device.type != 'cpu' 63 | init_seeds(2 + rank) 64 | with open(opt.data) as f: 65 | data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict 66 | is_coco = opt.data.endswith('coco.yaml') 67 | 68 | # Logging- Doing this before checking the dataset. Might update data_dict 69 | loggers = {'wandb': None} # loggers dict 70 | if rank in [-1, 0]: 71 | opt.hyp = hyp # add hyperparameters 72 | run_id = torch.load(weights, map_location=device).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None 73 | wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict) 74 | loggers['wandb'] = wandb_logger.wandb 75 | data_dict = wandb_logger.data_dict 76 | if wandb_logger.wandb: 77 | weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # WandbLogger might update weights, epochs if resuming 78 | 79 | nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes 80 | names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names 81 | assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check 82 | 83 | # Model 84 | pretrained = weights.endswith('.pt') 85 | if pretrained: 86 | with torch_distributed_zero_first(rank): 87 | attempt_download(weights) # download if not found locally 88 | ckpt = torch.load(weights, map_location=device) # load checkpoint 89 | model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 90 | exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys 91 | state_dict = ckpt['model'].float().state_dict() # to FP32 92 | state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect 93 | model.load_state_dict(state_dict, strict=False) # load 94 | logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report 95 | else: 96 | model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 97 | with torch_distributed_zero_first(rank): 98 | check_dataset(data_dict) # check 99 | train_path = data_dict['train'] 100 | test_path = data_dict['val'] 101 | 102 | # Freeze 103 | freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # parameter names to freeze (full or partial) 104 | for k, v in model.named_parameters(): 105 | v.requires_grad = True # train all layers 106 | if any(x in k for x in freeze): 107 | print('freezing %s' % k) 108 | v.requires_grad = False 109 | 110 | # Optimizer 111 | nbs = 64 # nominal batch size 112 | accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing 113 | hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay 114 | logger.info(f"Scaled weight_decay = {hyp['weight_decay']}") 115 | 116 | pg0, pg1, pg2 = [], [], [] # optimizer parameter groups 117 | for k, v in model.named_modules(): 118 | if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): 119 | pg2.append(v.bias) # biases 120 | if isinstance(v, nn.BatchNorm2d): 121 | pg0.append(v.weight) # no decay 122 | elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): 123 | pg1.append(v.weight) # apply decay 124 | if hasattr(v, 'im'): 125 | if hasattr(v.im, 'implicit'): 126 | pg0.append(v.im.implicit) 127 | else: 128 | for iv in v.im: 129 | pg0.append(iv.implicit) 130 | if hasattr(v, 'imc'): 131 | if hasattr(v.imc, 'implicit'): 132 | pg0.append(v.imc.implicit) 133 | else: 134 | for iv in v.imc: 135 | pg0.append(iv.implicit) 136 | if hasattr(v, 'imb'): 137 | if hasattr(v.imb, 'implicit'): 138 | pg0.append(v.imb.implicit) 139 | else: 140 | for iv in v.imb: 141 | pg0.append(iv.implicit) 142 | if hasattr(v, 'imo'): 143 | if hasattr(v.imo, 'implicit'): 144 | pg0.append(v.imo.implicit) 145 | else: 146 | for iv in v.imo: 147 | pg0.append(iv.implicit) 148 | if hasattr(v, 'ia'): 149 | if hasattr(v.ia, 'implicit'): 150 | pg0.append(v.ia.implicit) 151 | else: 152 | for iv in v.ia: 153 | pg0.append(iv.implicit) 154 | if hasattr(v, 'attn'): 155 | if hasattr(v.attn, 'logit_scale'): 156 | pg0.append(v.attn.logit_scale) 157 | if hasattr(v.attn, 'q_bias'): 158 | pg0.append(v.attn.q_bias) 159 | if hasattr(v.attn, 'v_bias'): 160 | pg0.append(v.attn.v_bias) 161 | if hasattr(v.attn, 'relative_position_bias_table'): 162 | pg0.append(v.attn.relative_position_bias_table) 163 | if hasattr(v, 'rbr_dense'): 164 | if hasattr(v.rbr_dense, 'weight_rbr_origin'): 165 | pg0.append(v.rbr_dense.weight_rbr_origin) 166 | if hasattr(v.rbr_dense, 'weight_rbr_avg_conv'): 167 | pg0.append(v.rbr_dense.weight_rbr_avg_conv) 168 | if hasattr(v.rbr_dense, 'weight_rbr_pfir_conv'): 169 | pg0.append(v.rbr_dense.weight_rbr_pfir_conv) 170 | if hasattr(v.rbr_dense, 'weight_rbr_1x1_kxk_idconv1'): 171 | pg0.append(v.rbr_dense.weight_rbr_1x1_kxk_idconv1) 172 | if hasattr(v.rbr_dense, 'weight_rbr_1x1_kxk_conv2'): 173 | pg0.append(v.rbr_dense.weight_rbr_1x1_kxk_conv2) 174 | if hasattr(v.rbr_dense, 'weight_rbr_gconv_dw'): 175 | pg0.append(v.rbr_dense.weight_rbr_gconv_dw) 176 | if hasattr(v.rbr_dense, 'weight_rbr_gconv_pw'): 177 | pg0.append(v.rbr_dense.weight_rbr_gconv_pw) 178 | if hasattr(v.rbr_dense, 'vector'): 179 | pg0.append(v.rbr_dense.vector) 180 | 181 | if opt.adam: 182 | optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum 183 | else: 184 | optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) 185 | 186 | optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay 187 | optimizer.add_param_group({'params': pg2}) # add pg2 (biases) 188 | logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0))) 189 | del pg0, pg1, pg2 190 | 191 | # Scheduler https://arxiv.org/pdf/1812.01187.pdf 192 | # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR 193 | if opt.linear_lr: 194 | lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear 195 | else: 196 | lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf'] 197 | scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) 198 | # plot_lr_scheduler(optimizer, scheduler, epochs) 199 | 200 | # EMA 201 | ema = ModelEMA(model) if rank in [-1, 0] else None 202 | 203 | # Resume 204 | start_epoch, best_fitness = 0, 0.0 205 | if pretrained: 206 | # Optimizer 207 | if ckpt['optimizer'] is not None: 208 | optimizer.load_state_dict(ckpt['optimizer']) 209 | best_fitness = ckpt['best_fitness'] 210 | 211 | # EMA 212 | if ema and ckpt.get('ema'): 213 | ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) 214 | ema.updates = ckpt['updates'] 215 | 216 | # Results 217 | if ckpt.get('training_results') is not None: 218 | results_file.write_text(ckpt['training_results']) # write results.txt 219 | 220 | # Epochs 221 | start_epoch = ckpt['epoch'] + 1 222 | if opt.resume: 223 | assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs) 224 | if epochs < start_epoch: 225 | logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' % 226 | (weights, ckpt['epoch'], epochs)) 227 | epochs += ckpt['epoch'] # finetune additional epochs 228 | 229 | del ckpt, state_dict 230 | 231 | # Image sizes 232 | gs = max(int(model.stride.max()), 32) # grid size (max stride) 233 | nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj']) 234 | imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples 235 | 236 | # DP mode 237 | if cuda and rank == -1 and torch.cuda.device_count() > 1: 238 | model = torch.nn.DataParallel(model) 239 | 240 | # SyncBatchNorm 241 | if opt.sync_bn and cuda and rank != -1: 242 | model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) 243 | logger.info('Using SyncBatchNorm()') 244 | 245 | # Trainloader 246 | dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, 247 | hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank, 248 | world_size=opt.world_size, workers=opt.workers, 249 | image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: ')) 250 | mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class 251 | nb = len(dataloader) # number of batches 252 | assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1) 253 | 254 | # Process 0 255 | if rank in [-1, 0]: 256 | testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader 257 | hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1, 258 | world_size=opt.world_size, workers=opt.workers, 259 | pad=0.5, prefix=colorstr('val: '))[0] 260 | 261 | if not opt.resume: 262 | labels = np.concatenate(dataset.labels, 0) 263 | c = torch.tensor(labels[:, 0]) # classes 264 | # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency 265 | # model._initialize_biases(cf.to(device)) 266 | if plots: 267 | #plot_labels(labels, names, save_dir, loggers) 268 | if tb_writer: 269 | tb_writer.add_histogram('classes', c, 0) 270 | 271 | # Anchors 272 | if not opt.noautoanchor: 273 | check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) 274 | model.half().float() # pre-reduce anchor precision 275 | 276 | # DDP mode 277 | if cuda and rank != -1: 278 | model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank, 279 | # nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698 280 | find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules())) 281 | 282 | # Model parameters 283 | hyp['box'] *= 3. / nl # scale to layers 284 | hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers 285 | hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers 286 | hyp['label_smoothing'] = opt.label_smoothing 287 | model.nc = nc # attach number of classes to model 288 | model.hyp = hyp # attach hyperparameters to model 289 | model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou) 290 | model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights 291 | model.names = names 292 | 293 | # Start training 294 | t0 = time.time() 295 | nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations) 296 | # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training 297 | maps = np.zeros(nc) # mAP per class 298 | results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) 299 | scheduler.last_epoch = start_epoch - 1 # do not move 300 | scaler = amp.GradScaler(enabled=cuda) 301 | compute_loss_ota = ComputeLossOTA(model) # init loss class 302 | compute_loss = ComputeLoss(model) # init loss class 303 | logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n' 304 | f'Using {dataloader.num_workers} dataloader workers\n' 305 | f'Logging results to {save_dir}\n' 306 | f'Starting training for {epochs} epochs...') 307 | torch.save(model, wdir / 'init.pt') 308 | for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ 309 | model.train() 310 | 311 | # Update image weights (optional) 312 | if opt.image_weights: 313 | # Generate indices 314 | if rank in [-1, 0]: 315 | cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights 316 | iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights 317 | dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx 318 | # Broadcast if DDP 319 | if rank != -1: 320 | indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int() 321 | dist.broadcast(indices, 0) 322 | if rank != 0: 323 | dataset.indices = indices.cpu().numpy() 324 | 325 | # Update mosaic border 326 | # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) 327 | # dataset.mosaic_border = [b - imgsz, -b] # height, width borders 328 | 329 | mloss = torch.zeros(4, device=device) # mean losses 330 | if rank != -1: 331 | dataloader.sampler.set_epoch(epoch) 332 | pbar = enumerate(dataloader) 333 | logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size')) 334 | if rank in [-1, 0]: 335 | pbar = tqdm(pbar, total=nb) # progress bar 336 | optimizer.zero_grad() 337 | for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- 338 | ni = i + nb * epoch # number integrated batches (since train start) 339 | imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0 340 | 341 | # Warmup 342 | if ni <= nw: 343 | xi = [0, nw] # x interp 344 | # model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou) 345 | accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round()) 346 | for j, x in enumerate(optimizer.param_groups): 347 | # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 348 | x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)]) 349 | if 'momentum' in x: 350 | x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) 351 | 352 | # Multi-scale 353 | if opt.multi_scale: 354 | sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size 355 | sf = sz / max(imgs.shape[2:]) # scale factor 356 | if sf != 1: 357 | ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) 358 | imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) 359 | 360 | # Forward 361 | with amp.autocast(enabled=cuda): 362 | pred = model(imgs) # forward 363 | if 'loss_ota' not in hyp or hyp['loss_ota'] == 1: 364 | loss, loss_items = compute_loss_ota(pred, targets.to(device), imgs) # loss scaled by batch_size 365 | else: 366 | loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size 367 | if rank != -1: 368 | loss *= opt.world_size # gradient averaged between devices in DDP mode 369 | if opt.quad: 370 | loss *= 4. 371 | 372 | # Backward 373 | scaler.scale(loss).backward() 374 | 375 | # Optimize 376 | if ni % accumulate == 0: 377 | scaler.step(optimizer) # optimizer.step 378 | scaler.update() 379 | optimizer.zero_grad() 380 | if ema: 381 | ema.update(model) 382 | 383 | # Print 384 | if rank in [-1, 0]: 385 | mloss = (mloss * i + loss_items) / (i + 1) # update mean losses 386 | mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB) 387 | s = ('%10s' * 2 + '%10.4g' * 6) % ( 388 | '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1]) 389 | pbar.set_description(s) 390 | 391 | # Plot 392 | if plots and ni < 10: 393 | f = save_dir / f'train_batch{ni}.jpg' # filename 394 | Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() 395 | # if tb_writer: 396 | # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) 397 | # tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph 398 | elif plots and ni == 10 and wandb_logger.wandb: 399 | wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in 400 | save_dir.glob('train*.jpg') if x.exists()]}) 401 | 402 | # end batch ------------------------------------------------------------------------------------------------ 403 | # end epoch ---------------------------------------------------------------------------------------------------- 404 | 405 | # Scheduler 406 | lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard 407 | scheduler.step() 408 | 409 | # DDP process 0 or single-GPU 410 | if rank in [-1, 0]: 411 | # mAP 412 | ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights']) 413 | final_epoch = epoch + 1 == epochs 414 | if not opt.notest or final_epoch: # Calculate mAP 415 | wandb_logger.current_epoch = epoch + 1 416 | results, maps, times = test.test(data_dict, 417 | batch_size=batch_size * 2, 418 | imgsz=imgsz_test, 419 | model=ema.ema, 420 | single_cls=opt.single_cls, 421 | dataloader=testloader, 422 | save_dir=save_dir, 423 | verbose=nc < 50 and final_epoch, 424 | plots=plots and final_epoch, 425 | wandb_logger=wandb_logger, 426 | compute_loss=compute_loss, 427 | is_coco=is_coco) 428 | 429 | # Write 430 | with open(results_file, 'a') as f: 431 | f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss 432 | if len(opt.name) and opt.bucket: 433 | os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name)) 434 | 435 | # Log 436 | tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss 437 | 'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 438 | 'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss 439 | 'x/lr0', 'x/lr1', 'x/lr2'] # params 440 | for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags): 441 | if tb_writer: 442 | tb_writer.add_scalar(tag, x, epoch) # tensorboard 443 | if wandb_logger.wandb: 444 | wandb_logger.log({tag: x}) # W&B 445 | 446 | # Update best mAP 447 | fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] 448 | if fi > best_fitness: 449 | best_fitness = fi 450 | wandb_logger.end_epoch(best_result=best_fitness == fi) 451 | 452 | # Save model 453 | if (not opt.nosave) or (final_epoch and not opt.evolve): # if save 454 | ckpt = {'epoch': epoch, 455 | 'best_fitness': best_fitness, 456 | 'training_results': results_file.read_text(), 457 | 'model': deepcopy(model.module if is_parallel(model) else model).half(), 458 | 'ema': deepcopy(ema.ema).half(), 459 | 'updates': ema.updates, 460 | 'optimizer': optimizer.state_dict(), 461 | 'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None} 462 | 463 | # Save last, best and delete 464 | torch.save(ckpt, last) 465 | if best_fitness == fi: 466 | torch.save(ckpt, best) 467 | if (best_fitness == fi) and (epoch >= 200): 468 | torch.save(ckpt, wdir / 'best_{:03d}.pt'.format(epoch)) 469 | if epoch == 0: 470 | torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) 471 | elif ((epoch+1) % 25) == 0: 472 | torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) 473 | elif epoch >= (epochs-5): 474 | torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) 475 | if wandb_logger.wandb: 476 | if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1: 477 | wandb_logger.log_model( 478 | last.parent, opt, epoch, fi, best_model=best_fitness == fi) 479 | del ckpt 480 | 481 | # end epoch ---------------------------------------------------------------------------------------------------- 482 | # end training 483 | if rank in [-1, 0]: 484 | # Plots 485 | if plots: 486 | plot_results(save_dir=save_dir) # save as results.png 487 | if wandb_logger.wandb: 488 | files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]] 489 | wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files 490 | if (save_dir / f).exists()]}) 491 | # Test best.pt 492 | logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600)) 493 | if opt.data.endswith('coco.yaml') and nc == 80: # if COCO 494 | for m in (last, best) if best.exists() else (last): # speed, mAP tests 495 | results, _, _ = test.test(opt.data, 496 | batch_size=batch_size * 2, 497 | imgsz=imgsz_test, 498 | conf_thres=0.001, 499 | iou_thres=0.7, 500 | model=attempt_load(m, device).half(), 501 | single_cls=opt.single_cls, 502 | dataloader=testloader, 503 | save_dir=save_dir, 504 | save_json=True, 505 | plots=False, 506 | is_coco=is_coco) 507 | 508 | # Strip optimizers 509 | final = best if best.exists() else last # final model 510 | for f in last, best: 511 | if f.exists(): 512 | strip_optimizer(f) # strip optimizers 513 | if opt.bucket: 514 | os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload 515 | if wandb_logger.wandb and not opt.evolve: # Log the stripped model 516 | wandb_logger.wandb.log_artifact(str(final), type='model', 517 | name='run_' + wandb_logger.wandb_run.id + '_model', 518 | aliases=['last', 'best', 'stripped']) 519 | wandb_logger.finish_run() 520 | else: 521 | dist.destroy_process_group() 522 | torch.cuda.empty_cache() 523 | return results 524 | 525 | 526 | if __name__ == '__main__': 527 | parser = argparse.ArgumentParser() 528 | parser.add_argument('--weights', type=str, default='yolov7.pt', help='initial weights path') 529 | parser.add_argument('--cfg', type=str, default='cfg/training/yolov7.yaml', help='model.yaml path') 530 | parser.add_argument('--data', type=str, default='data/coco.yaml', help='data.yaml path') 531 | parser.add_argument('--hyp', type=str, default='data/hyp.scratch.p5.yaml', help='hyperparameters path') 532 | parser.add_argument('--epochs', type=int, default=300) 533 | parser.add_argument('--batch-size', type=int, default=4, help='total batch size for all GPUs') 534 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') 535 | parser.add_argument('--rect', action='store_true', help='rectangular training') 536 | parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') 537 | parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') 538 | parser.add_argument('--notest', action='store_true', help='only test final epoch') 539 | parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check') 540 | parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') 541 | parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') 542 | parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') 543 | parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') 544 | parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 545 | parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') 546 | parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class') 547 | parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') 548 | parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') 549 | parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') 550 | parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers') 551 | parser.add_argument('--project', default='runs/train', help='save to project/name') 552 | parser.add_argument('--entity', default=None, help='W&B entity') 553 | parser.add_argument('--name', default='exp', help='save to project/name') 554 | parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') 555 | parser.add_argument('--quad', action='store_true', help='quad dataloader') 556 | parser.add_argument('--linear-lr', action='store_true', help='linear LR') 557 | parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon') 558 | parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table') 559 | parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B') 560 | parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch') 561 | parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used') 562 | parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone of yolov7=50, first3=0 1 2') 563 | opt = parser.parse_args() 564 | 565 | # Set DDP variables 566 | opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1 567 | opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1 568 | set_logging(opt.global_rank) 569 | #if opt.global_rank in [-1, 0]: 570 | # check_git_status() 571 | # check_requirements() 572 | 573 | # Resume 574 | wandb_run = check_wandb_resume(opt) 575 | if opt.resume and not wandb_run: # resume an interrupted run 576 | ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path 577 | assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' 578 | apriori = opt.global_rank, opt.local_rank 579 | with open(Path(ckpt).parent.parent / 'opt.yaml') as f: 580 | opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader)) # replace 581 | opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori # reinstate 582 | logger.info('Resuming training from %s' % ckpt) 583 | else: 584 | # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml') 585 | opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp) # check files 586 | assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified' 587 | opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test) 588 | opt.name = 'evolve' if opt.evolve else opt.name 589 | opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve) # increment run 590 | 591 | # DDP mode 592 | opt.total_batch_size = opt.batch_size 593 | device = select_device(opt.device, batch_size=opt.batch_size) 594 | if opt.local_rank != -1: 595 | assert torch.cuda.device_count() > opt.local_rank 596 | torch.cuda.set_device(opt.local_rank) 597 | device = torch.device('cuda', opt.local_rank) 598 | dist.init_process_group(backend='nccl', init_method='env://') # distributed backend 599 | assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count' 600 | opt.batch_size = opt.total_batch_size // opt.world_size 601 | 602 | # Hyperparameters 603 | with open(opt.hyp) as f: 604 | hyp = yaml.load(f, Loader=yaml.SafeLoader) # load hyps 605 | 606 | # Train 607 | logger.info(opt) 608 | if not opt.evolve: 609 | tb_writer = None # init loggers 610 | if opt.global_rank in [-1, 0]: 611 | prefix = colorstr('tensorboard: ') 612 | logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/") 613 | tb_writer = SummaryWriter(opt.save_dir) # Tensorboard 614 | train(hyp, opt, device, tb_writer) 615 | 616 | # Evolve hyperparameters (optional) 617 | else: 618 | # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) 619 | meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) 620 | 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) 621 | 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 622 | 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay 623 | 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok) 624 | 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum 625 | 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr 626 | 'box': (1, 0.02, 0.2), # box loss gain 627 | 'cls': (1, 0.2, 4.0), # cls loss gain 628 | 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight 629 | 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) 630 | 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight 631 | 'iou_t': (0, 0.1, 0.7), # IoU training threshold 632 | 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold 633 | 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore) 634 | 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) 635 | 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) 636 | 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) 637 | 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) 638 | 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) 639 | 'translate': (1, 0.0, 0.9), # image translation (+/- fraction) 640 | 'scale': (1, 0.0, 0.9), # image scale (+/- gain) 641 | 'shear': (1, 0.0, 10.0), # image shear (+/- deg) 642 | 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 643 | 'flipud': (1, 0.0, 1.0), # image flip up-down (probability) 644 | 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) 645 | 'mosaic': (1, 0.0, 1.0), # image mixup (probability) 646 | 'mixup': (1, 0.0, 1.0), # image mixup (probability) 647 | 'copy_paste': (1, 0.0, 1.0), # segment copy-paste (probability) 648 | 'paste_in': (1, 0.0, 1.0)} # segment copy-paste (probability) 649 | 650 | with open(opt.hyp, errors='ignore') as f: 651 | hyp = yaml.safe_load(f) # load hyps dict 652 | if 'anchors' not in hyp: # anchors commented in hyp.yaml 653 | hyp['anchors'] = 3 654 | 655 | assert opt.local_rank == -1, 'DDP mode not implemented for --evolve' 656 | opt.notest, opt.nosave = True, True # only test/save final epoch 657 | # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices 658 | yaml_file = Path(opt.save_dir) / 'hyp_evolved.yaml' # save best result here 659 | if opt.bucket: 660 | os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket) # download evolve.txt if exists 661 | 662 | for _ in range(300): # generations to evolve 663 | if Path('evolve.txt').exists(): # if evolve.txt exists: select best hyps and mutate 664 | # Select parent(s) 665 | parent = 'single' # parent selection method: 'single' or 'weighted' 666 | x = np.loadtxt('evolve.txt', ndmin=2) 667 | n = min(5, len(x)) # number of previous results to consider 668 | x = x[np.argsort(-fitness(x))][:n] # top n mutations 669 | w = fitness(x) - fitness(x).min() # weights 670 | if parent == 'single' or len(x) == 1: 671 | # x = x[random.randint(0, n - 1)] # random selection 672 | x = x[random.choices(range(n), weights=w)[0]] # weighted selection 673 | elif parent == 'weighted': 674 | x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination 675 | 676 | # Mutate 677 | mp, s = 0.8, 0.2 # mutation probability, sigma 678 | npr = np.random 679 | npr.seed(int(time.time())) 680 | g = np.array([x[0] for x in meta.values()]) # gains 0-1 681 | ng = len(meta) 682 | v = np.ones(ng) 683 | while all(v == 1): # mutate until a change occurs (prevent duplicates) 684 | v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0) 685 | for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300) 686 | hyp[k] = float(x[i + 7] * v[i]) # mutate 687 | 688 | # Constrain to limits 689 | for k, v in meta.items(): 690 | hyp[k] = max(hyp[k], v[1]) # lower limit 691 | hyp[k] = min(hyp[k], v[2]) # upper limit 692 | hyp[k] = round(hyp[k], 5) # significant digits 693 | 694 | # Train mutation 695 | results = train(hyp.copy(), opt, device) 696 | 697 | # Write mutation results 698 | print_mutation(hyp.copy(), results, yaml_file, opt.bucket) 699 | 700 | # Plot results 701 | plot_evolution(yaml_file) 702 | print(f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n' 703 | f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}') 704 | --------------------------------------------------------------------------------