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