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