经验首页 前端设计 程序设计 Java相关 移动开发 数据库/运维 软件/图像 大数据/云计算 其他经验
当前位置：技术经验 » 大数据/云/AI » 人工智能基础 » 查看文章
训练一个目标检测模型
来源：cnblogs　　作者：zylyehuo　　时间：2023/12/15 8:50:42　　对本文有异议

博客地址：https://www.cnblogs.com/zylyehuo/

（一）识别背景/目的

第十八届全国大学生智能汽车竞赛室外 ROS 无人车赛（高教组）

无人车在室外运行中， 需要探索未知环境， 识别障碍物， 停车标志牌、红绿灯等标志物。

比赛场地为不规则环形场地， 由红蓝两色锥桶搭建而成， 整体赛道由直线区域、 "S"弯、 直角区域、 圆形区域等部分元素或全部元素构成

（二）识别/运行场地

① 一食堂二楼

② 室外网球场地

（三）实现效果

（四）技术栈

  • 识别模型：yolov5
  • 标注工具：labelmaster
  • 运行环境：Ubuntu20.04
  • 加速方式：使用onnx及tensorrt将模型进行推理加速

（五）识别类别

  • 0: red 红色锥桶
  • 1: blue 蓝色锥桶
  • 2: stop 红绿灯（红灯）
  • 3: wait 停车牌

（六）yolov5目标检测模型及其环境配置过程

第一步：下载yolov5源码

https://github.com/ultralytics/yolov5

第二步：解压源码压缩包

第三步：在代码编辑器 pycharm/vscode 中打开源码文件夹（配置完成）

（七）训练数据、测试数据采集

第一步：打开 ROS 智能车摄像头

终端输入 cheese

第二步：点击拍摄按键，采集数据集图片

第三步：将无人车上的照片拷贝到电脑上，为数据标注做准备

（八）训练数据、测试数据标注、整理

第一步：安装 labelmaster 库

  1. pip install labelImg

第二步：启动 labelmaster

  1. labelImg

第三步：打开采集图片的保存路径

第四步：鼠标右键图片，创建区块标注

左侧选择创建区块

鼠标移至目标的左上角

点击鼠标并拉直右下角

在出现的框里面选择自己标注目标的分类

键入新的分类，则会自动生成一个新的分类

这边选择wait的红绿灯分类

左键单击 ok 键

完成单个目标的标注

其他类别同样操作，只是分类时选择不同分类

依次按照识别类别对其余区块进行标注

例如下图选择red的红色锥桶分类

全部完成后左上角点击“改变保存目录”

选择保存的目标文件夹

第五步：保存标注参数，保存为 .txt 文件后缀

点击保存下方的按键，改变保存的文件格式，保存为 .txt 文件后缀

“yolo”对应的是“txt”文件

“PascalVOC”对应的是“xml”文件

“CreatelML”对应的是“json”文件

最后单击保存完成标注

此时，保存的文件夹中会自动生成一个 class.txt 文件

（九）模型训练过程

为了减小最终模型保存的大小，保证无人车整体运行的流程性，我们选用 yolov5n.yaml 的参数作为样本

源码：yolov5n.yaml

  1. # Parameters
  2. nc: 80  # number of classes
  3. depth_multiple: 0.33  # model depth multiple
  4. width_multiple: 0.25  # layer channel multiple
  5. anchors:
  6.   - [10,13, 16,30, 33,23]  # P3/8
  7.   - [30,61, 62,45, 59,119]  # P4/16
  8.   - [116,90, 156,198, 373,326]  # P5/32
  10. # YOLOv5 v6.0 backbone
  11. backbone:
  12.   # [from, number, module, args]
  13.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
  14.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
  15.    [-1, 3, C3, [128]],
  16.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
  17.    [-1, 6, C3, [256]],
  18.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
  19.    [-1, 9, C3, [512]],
  20.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
  21.    [-1, 3, C3, [1024]],
  22.    [-1, 1, SPPF, [1024, 5]],  # 9
  23.   ]
  25. # YOLOv5 v6.0 head
  26. head:
  27.   [[-1, 1, Conv, [512, 1, 1]],
  28.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  29.    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
  30.    [-1, 3, C3, [512, False]],  # 13
  32.    [-1, 1, Conv, [256, 1, 1]],
  33.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  34.    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
  35.    [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
  37.    [-1, 1, Conv, [256, 3, 2]],
  38.    [[-1, 14], 1, Concat, [1]],  # cat head P4
  39.    [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
  41.    [-1, 1, Conv, [512, 3, 2]],
  42.    [[-1, 10], 1, Concat, [1]],  # cat head P5
  43.    [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
  45.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  46.   ]

第一步：配置识别模型参数

① IP_model.yaml（在 yolov5n.yaml 基础上进行修改）

  1. # Parameters
  2. nc: 4  # number of classes
  3. depth_multiple: 0.33  # model depth multiple
  4. width_multiple: 0.25  # layer channel multiple
  5. anchors:
  6.   - [10,13, 16,30, 33,23]  # P3/8
  7.   - [30,61, 62,45, 59,119]  # P4/16
  8.   - [116,90, 156,198, 373,326]  # P5/32
  10. # YOLOv5 v6.0 backbone
  11. backbone:
  12.   # [from, number, module, args]
  13.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
  14.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
  15.    [-1, 3, C3, [128]],
  16.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
  17.    [-1, 6, C3, [256]],
  18.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
  19.    [-1, 9, C3, [512]],
  20.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
  21.    [-1, 3, C3, [1024]],
  22.    [-1, 1, SPPF, [1024, 5]],  # 9
  23.   ]
  25. # YOLOv5 v6.0 head
  26. head:
  27.   [[-1, 1, Conv, [512, 1, 1]],
  28.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  29.    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
  30.    [-1, 3, C3, [512, False]],  # 13
  32.    [-1, 1, Conv, [256, 1, 1]],
  33.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  34.    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
  35.    [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
  37.    [-1, 1, Conv, [256, 3, 2]],
  38.    [[-1, 14], 1, Concat, [1]],  # cat head P4
  39.    [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
  41.    [-1, 1, Conv, [512, 3, 2]],
  42.    [[-1, 10], 1, Concat, [1]],  # cat head P5
  43.    [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
  45.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  46.   ]

② IP_parameter.yaml（模型训练以及模型保存路径等参数设置）

  1. path: D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\V2\datasets  # dataset root dir
  2. train: images/train  # train images (relative to 'path') 128 images
  3. val: images/train  # val images (relative to 'path') 128 images
  4. test:  # test images (optional)
  6. # Classes
  7. names:
  8.   0: red
  9.   1: blue
  10.   2: stop
  11.   3: wait
  14. # Download script/URL (optional)
  15. download: https://ultralytics.com/assets/coco128.zip

第二步：编写模型训练代码

train.py

  1. import argparse
  2. import math
  3. import os
  4. os.environ["GIT_PYTHON_REFRESH"] = "quiet"
  5. import random
  6. import subprocess
  7. import sys
  8. import time
  9. from copy import deepcopy
  10. from datetime import datetime
  11. from pathlib import Path
  13. import numpy as np
  14. import torch
  15. import torch.distributed as dist
  16. import torch.nn as nn
  17. import yaml
  18. from torch.optim import lr_scheduler
  19. from tqdm import tqdm
  21. FILE = Path(__file__).resolve()
  22. ROOT = FILE.parents[0]  # YOLOv5 root directory
  23. if str(ROOT) not in sys.path:
  24.     sys.path.append(str(ROOT))  # add ROOT to PATH
  25. ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
  27. import val as validate  # for end-of-epoch mAP
  28. from models.experimental import attempt_load
  29. from models.yolo import Model
  30. from utils.autoanchor import check_anchors
  31. from utils.autobatch import check_train_batch_size
  32. from utils.callbacks import Callbacks
  33. from utils.dataloaders import create_dataloader
  34. from utils.downloads import attempt_download, is_url
  35. from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
  36.                            check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,
  37.                            get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
  38.                            labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer,
  39.                            yaml_save)
  40. from utils.loggers import Loggers
  41. from utils.loggers.comet.comet_utils import check_comet_resume
  42. from utils.loss import ComputeLoss
  43. from utils.metrics import fitness
  44. from utils.plots import plot_evolve
  45. from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
  46.                                smart_resume, torch_distributed_zero_first)
  48. LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
  49. RANK = int(os.getenv('RANK', -1))
  50. WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
  51. GIT_INFO = check_git_info()
  54. def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictionary
  55.     save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze =         Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg,         opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
  56.     callbacks.run('on_pretrain_routine_start')
  58.     # Directories
  59.     w = save_dir / 'weights'  # weights dir
  60.     (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
  61.     last, best = w / 'last.pt', w / 'best.pt'
  63.     # Hyperparameters
  64.     if isinstance(hyp, str):
  65.         with open(hyp, errors='ignore') as f:
  66.             hyp = yaml.safe_load(f)  # load hyps dict
  67.     LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
  68.     opt.hyp = hyp.copy()  # for saving hyps to checkpoints
  70.     # Save run settings
  71.     if not evolve:
  72.         yaml_save(save_dir / 'hyp.yaml', hyp)
  73.         yaml_save(save_dir / 'opt.yaml', vars(opt))
  75.     # Loggers
  76.     data_dict = None
  77.     if RANK in {-1, 0}:
  78.         loggers = Loggers(save_dir, weights, opt, hyp, LOGGER)  # loggers instance
  80.         # Register actions
  81.         for k in methods(loggers):
  82.             callbacks.register_action(k, callback=getattr(loggers, k))
  84.         # Process custom dataset artifact link
  85.         data_dict = loggers.remote_dataset
  86.         if resume:  # If resuming runs from remote artifact
  87.             weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
  89.     # Config
  90.     plots = not evolve and not opt.noplots  # create plots
  91.     cuda = device.type != 'cpu'
  92.     init_seeds(opt.seed + 1 + RANK, deterministic=True)
  93.     with torch_distributed_zero_first(LOCAL_RANK):
  94.         data_dict = data_dict or check_dataset(data)  # check if None
  95.     train_path, val_path = data_dict['train'], data_dict['val']
  96.     nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
  97.     names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
  98.     is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')  # COCO dataset
  100.     # Model
  101.     check_suffix(weights, '.pt')  # check weights
  102.     pretrained = weights.endswith('.pt')
  103.     if pretrained:
  104.         with torch_distributed_zero_first(LOCAL_RANK):
  105.             weights = attempt_download(weights)  # download if not found locally
  106.         ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
  107.         model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  108.         exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []  # exclude keys
  109.         csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
  110.         csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
  111.         model.load_state_dict(csd, strict=False)  # load
  112.         LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  # report
  113.     else:
  114.         model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  115.     amp = check_amp(model)  # check AMP
  117.     # Freeze
  118.     freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
  119.     for k, v in model.named_parameters():
  120.         v.requires_grad = True  # train all layers
  121.         # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
  122.         if any(x in k for x in freeze):
  123.             LOGGER.info(f'freezing {k}')
  124.             v.requires_grad = False
  126.     # Image size
  127.     gs = max(int(model.stride.max()), 32)  # grid size (max stride)
  128.     imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)  # verify imgsz is gs-multiple
  130.     # Batch size
  131.     if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
  132.         batch_size = check_train_batch_size(model, imgsz, amp)
  133.         loggers.on_params_update({'batch_size': batch_size})
  135.     # Optimizer
  136.     nbs = 64  # nominal batch size
  137.     accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
  138.     hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
  139.     optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
  141.     # Scheduler
  142.     if opt.cos_lr:
  143.         lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
  144.     else:
  145.         lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
  146.     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)
  148.     # EMA
  149.     ema = ModelEMA(model) if RANK in {-1, 0} else None
  151.     # Resume
  152.     best_fitness, start_epoch = 0.0, 0
  153.     if pretrained:
  154.         if resume:
  155.             best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
  156.         del ckpt, csd
  158.     # DP mode
  159.     if cuda and RANK == -1 and torch.cuda.device_count() > 1:
  160.         LOGGER.warning(
  161.             'WARNING ?? DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
  162.             'See Multi-GPU Tutorial at https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training to get started.'
  163.         )
  164.         model = torch.nn.DataParallel(model)
  166.     # SyncBatchNorm
  167.     if opt.sync_bn and cuda and RANK != -1:
  168.         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
  169.         LOGGER.info('Using SyncBatchNorm()')
  171.     # Trainloader
  172.     train_loader, dataset = create_dataloader(train_path,
  173.                                               imgsz,
  174.                                               batch_size // WORLD_SIZE,
  175.                                               gs,
  176.                                               single_cls,
  177.                                               hyp=hyp,
  178.                                               augment=True,
  179.                                               cache=None if opt.cache == 'val' else opt.cache,
  180.                                               rect=opt.rect,
  181.                                               rank=LOCAL_RANK,
  182.                                               workers=workers,
  183.                                               image_weights=opt.image_weights,
  184.                                               quad=opt.quad,
  185.                                               prefix=colorstr('train: '),
  186.                                               shuffle=True,
  187.                                               seed=opt.seed)
  188.     labels = np.concatenate(dataset.labels, 0)
  189.     mlc = int(labels[:, 0].max())  # max label class
  190.     assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
  192.     # Process 0
  193.     if RANK in {-1, 0}:
  194.         val_loader = create_dataloader(val_path,
  195.                                        imgsz,
  196.                                        batch_size // WORLD_SIZE * 2,
  197.                                        gs,
  198.                                        single_cls,
  199.                                        hyp=hyp,
  200.                                        cache=None if noval else opt.cache,
  201.                                        rect=True,
  202.                                        rank=-1,
  203.                                        workers=workers * 2,
  204.                                        pad=0.5,
  205.                                        prefix=colorstr('val: '))[0]
  207.         if not resume:
  208.             if not opt.noautoanchor:
  209.                 check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)  # run AutoAnchor
  210.             model.half().float()  # pre-reduce anchor precision
  212.         callbacks.run('on_pretrain_routine_end', labels, names)
  214.     # DDP mode
  215.     if cuda and RANK != -1:
  216.         model = smart_DDP(model)
  218.     # Model attributes
  219.     nl = de_parallel(model).model[-1].nl  # number of detection layers (to scale hyps)
  220.     hyp['box'] *= 3 / nl  # scale to layers
  221.     hyp['cls'] *= nc / 80 * 3 / nl  # scale to classes and layers
  222.     hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
  223.     hyp['label_smoothing'] = opt.label_smoothing
  224.     model.nc = nc  # attach number of classes to model
  225.     model.hyp = hyp  # attach hyperparameters to model
  226.     model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
  227.     model.names = names
  229.     # Start training
  230.     t0 = time.time()
  231.     nb = len(train_loader)  # number of batches
  232.     nw = max(round(hyp['warmup_epochs'] * nb), 100)  # number of warmup iterations, max(3 epochs, 100 iterations)
  233.     # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
  234.     last_opt_step = -1
  235.     maps = np.zeros(nc)  # mAP per class
  236.     results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
  237.     scheduler.last_epoch = start_epoch - 1  # do not move
  238.     scaler = torch.cuda.amp.GradScaler(enabled=amp)
  239.     stopper, stop = EarlyStopping(patience=opt.patience), False
  240.     compute_loss = ComputeLoss(model)  # init loss class
  241.     callbacks.run('on_train_start')
  242.     LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
  243.                 f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
  244.                 f"Logging results to {colorstr('bold', save_dir)}\n"
  245.                 f'Starting training for {epochs} epochs...')
  246.     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
  247.         callbacks.run('on_train_epoch_start')
  248.         model.train()
  250.         # Update image weights (optional, single-GPU only)
  251.         if opt.image_weights:
  252.             cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights
  253.             iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
  254.             dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
  256.         # Update mosaic border (optional)
  257.         # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
  258.         # dataset.mosaic_border = [b - imgsz, -b]  # height, width borders
  260.         mloss = torch.zeros(3, device=device)  # mean losses
  261.         if RANK != -1:
  262.             train_loader.sampler.set_epoch(epoch)
  263.         pbar = enumerate(train_loader)
  264.         LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
  265.         if RANK in {-1, 0}:
  266.             pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT)  # progress bar
  267.         optimizer.zero_grad()
  268.         for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
  269.             callbacks.run('on_train_batch_start')
  270.             ni = i + nb * epoch  # number integrated batches (since train start)
  271.             imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0
  273.             # Warmup
  274.             if ni <= nw:
  275.                 xi = [0, nw]  # x interp
  276.                 # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
  277.                 accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
  278.                 for j, x in enumerate(optimizer.param_groups):
  279.                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
  280.                     x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
  281.                     if 'momentum' in x:
  282.                         x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
  284.             # Multi-scale
  285.             if opt.multi_scale:
  286.                 sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs  # size
  287.                 sf = sz / max(imgs.shape[2:])  # scale factor
  288.                 if sf != 1:
  289.                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
  290.                     imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
  292.             # Forward
  293.             with torch.cuda.amp.autocast(amp):
  294.                 pred = model(imgs)  # forward
  295.                 loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
  296.                 if RANK != -1:
  297.                     loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
  298.                 if opt.quad:
  299.                     loss *= 4.
  301.             # Backward
  302.             scaler.scale(loss).backward()
  304.             # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
  305.             if ni - last_opt_step >= accumulate:
  306.                 scaler.unscale_(optimizer)  # unscale gradients
  307.                 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)  # clip gradients
  308.                 scaler.step(optimizer)  # optimizer.step
  309.                 scaler.update()
  310.                 optimizer.zero_grad()
  311.                 if ema:
  312.                     ema.update(model)
  313.                 last_opt_step = ni
  315.             # Log
  316.             if RANK in {-1, 0}:
  317.                 mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
  318.                 mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'  # (GB)
  319.                 pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
  320.                                      (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
  321.                 callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
  322.                 if callbacks.stop_training:
  323.                     return
  324.             # end batch ------------------------------------------------------------------------------------------------
  326.         # Scheduler
  327.         lr = [x['lr'] for x in optimizer.param_groups]  # for loggers
  328.         scheduler.step()
  330.         if RANK in {-1, 0}:
  331.             # mAP
  332.             callbacks.run('on_train_epoch_end', epoch=epoch)
  333.             ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
  334.             final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
  335.             if not noval or final_epoch:  # Calculate mAP
  336.                 results, maps, _ = validate.run(data_dict,
  337.                                                 batch_size=batch_size // WORLD_SIZE * 2,
  338.                                                 imgsz=imgsz,
  339.                                                 half=amp,
  340.                                                 model=ema.ema,
  341.                                                 single_cls=single_cls,
  342.                                                 dataloader=val_loader,
  343.                                                 save_dir=save_dir,
  344.                                                 plots=False,
  345.                                                 callbacks=callbacks,
  346.                                                 compute_loss=compute_loss)
  348.             # Update best mAP
  349.             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
  350.             stop = stopper(epoch=epoch, fitness=fi)  # early stop check
  351.             if fi > best_fitness:
  352.                 best_fitness = fi
  353.             log_vals = list(mloss) + list(results) + lr
  354.             callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
  356.             # Save model
  357.             if (not nosave) or (final_epoch and not evolve):  # if save
  358.                 ckpt = {
  359.                     'epoch': epoch,
  360.                     'best_fitness': best_fitness,
  361.                     'model': deepcopy(de_parallel(model)).half(),
  362.                     'ema': deepcopy(ema.ema).half(),
  363.                     'updates': ema.updates,
  364.                     'optimizer': optimizer.state_dict(),
  365.                     'opt': vars(opt),
  366.                     'git': GIT_INFO,  # {remote, branch, commit} if a git repo
  367.                     'date': datetime.now().isoformat()}
  369.                 # Save last, best and delete
  370.                 torch.save(ckpt, last)
  371.                 if best_fitness == fi:
  372.                     torch.save(ckpt, best)
  373.                 if opt.save_period > 0 and epoch % opt.save_period == 0:
  374.                     torch.save(ckpt, w / f'epoch{epoch}.pt')
  375.                 del ckpt
  376.                 callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
  378.         # EarlyStopping
  379.         if RANK != -1:  # if DDP training
  380.             broadcast_list = [stop if RANK == 0 else None]
  381.             dist.broadcast_object_list(broadcast_list, 0)  # broadcast 'stop' to all ranks
  382.             if RANK != 0:
  383.                 stop = broadcast_list[0]
  384.         if stop:
  385.             break  # must break all DDP ranks
  387.         # end epoch ----------------------------------------------------------------------------------------------------
  388.     # end training -----------------------------------------------------------------------------------------------------
  389.     if RANK in {-1, 0}:
  390.         LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
  391.         for f in last, best:
  392.             if f.exists():
  393.                 strip_optimizer(f)  # strip optimizers
  394.                 if f is best:
  395.                     LOGGER.info(f'\nValidating {f}...')
  396.                     results, _, _ = validate.run(
  397.                         data_dict,
  398.                         batch_size=batch_size // WORLD_SIZE * 2,
  399.                         imgsz=imgsz,
  400.                         model=attempt_load(f, device).half(),
  401.                         iou_thres=0.65 if is_coco else 0.60,  # best pycocotools at iou 0.65
  402.                         single_cls=single_cls,
  403.                         dataloader=val_loader,
  404.                         save_dir=save_dir,
  405.                         save_json=is_coco,
  406.                         verbose=True,
  407.                         plots=plots,
  408.                         callbacks=callbacks,
  409.                         compute_loss=compute_loss)  # val best model with plots
  410.                     if is_coco:
  411.                         callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
  413.         callbacks.run('on_train_end', last, best, epoch, results)
  415.     torch.cuda.empty_cache()
  416.     return results
  419. def parse_opt(known=False):
  420.     parser = argparse.ArgumentParser()
  421.     parser.add_argument('--weights', type=str, default=ROOT / r'yolov5n.pt', help='initial weights path')
  422.     parser.add_argument('--cfg', type=str, default=ROOT / 'D:/BaiduNetdiskWorkspace/licenses/machine_learning/yolov5/yolov5-master/V2/IP_model.yaml', help='model.yaml path')
  423.     parser.add_argument('--data', type=str, default=ROOT / 'D:/BaiduNetdiskWorkspace/licenses/machine_learning/yolov5/yolov5-master/V2/IP_parameter.yaml', help='dataset.yaml path')
  424.     parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
  425.     parser.add_argument('--epochs', type=int, default=300, help='total training epochs')
  426.     parser.add_argument('--batch-size', type=int, default=-1, help='total batch size for all GPUs, -1 for autobatch')
  427.     parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=256, help='train, val image size (pixels)')
  428.     parser.add_argument('--rect', action='store_true', help='rectangular training')
  429.     parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
  430.     parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
  431.     parser.add_argument('--noval', action='store_true', help='only validate final epoch')
  432.     parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
  433.     parser.add_argument('--noplots', action='store_true', help='save no plot files')
  434.     parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
  435.     parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
  436.     parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')
  437.     parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
  438.     parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  439.     parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
  440.     parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
  441.     parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
  442.     parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
  443.     parser.add_argument('--workers', type=int, default=6, help='max dataloader workers (per RANK in DDP mode)')
  444.     parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
  445.     parser.add_argument('--name', default='exp', help='save to project/name')
  446.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  447.     parser.add_argument('--quad', action='store_true', help='quad dataloader')
  448.     parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
  449.     parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
  450.     parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
  451.     parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
  452.     parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
  453.     parser.add_argument('--seed', type=int, default=0, help='Global training seed')
  454.     parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
  456.     # Logger arguments
  457.     parser.add_argument('--entity', default=None, help='Entity')
  458.     parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
  459.     parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
  460.     parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
  462.     return parser.parse_known_args()[0] if known else parser.parse_args()
  465. def main(opt, callbacks=Callbacks()):
  466.     # Checks
  467.     if RANK in {-1, 0}:
  468.         print_args(vars(opt))
  469.         check_git_status()
  470.         check_requirements()
  472.     # Resume (from specified or most recent last.pt)
  473.     if opt.resume and not check_comet_resume(opt) and not opt.evolve:
  474.         last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
  475.         opt_yaml = last.parent.parent / 'opt.yaml'  # train options yaml
  476.         opt_data = opt.data  # original dataset
  477.         if opt_yaml.is_file():
  478.             with open(opt_yaml, errors='ignore') as f:
  479.                 d = yaml.safe_load(f)
  480.         else:
  481.             d = torch.load(last, map_location='cpu')['opt']
  482.         opt = argparse.Namespace(**d)  # replace
  483.         opt.cfg, opt.weights, opt.resume = '', str(last), True  # reinstate
  484.         if is_url(opt_data):
  485.             opt.data = check_file(opt_data)  # avoid HUB resume auth timeout
  486.     else:
  487.         opt.data, opt.cfg, opt.hyp, opt.weights, opt.project =             check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project)  # checks
  488.         assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
  489.         if opt.evolve:
  490.             if opt.project == str(ROOT / 'runs/train'):  # if default project name, rename to runs/evolve
  491.                 opt.project = str(ROOT / 'runs/evolve')
  492.             opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resume
  493.         if opt.name == 'cfg':
  494.             opt.name = Path(opt.cfg).stem  # use model.yaml as name
  495.         opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
  497.     # DDP mode
  498.     device = select_device(opt.device, batch_size=opt.batch_size)
  499.     if LOCAL_RANK != -1:
  500.         msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
  501.         assert not opt.image_weights, f'--image-weights {msg}'
  502.         assert not opt.evolve, f'--evolve {msg}'
  503.         assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
  504.         assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
  505.         assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
  506.         torch.cuda.set_device(LOCAL_RANK)
  507.         device = torch.device('cuda', LOCAL_RANK)
  508.         dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo')
  510.     # Train
  511.     if not opt.evolve:
  512.         train(opt.hyp, opt, device, callbacks)
  514.     # Evolve hyperparameters (optional)
  515.     else:
  516.         # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
  517.         meta = {
  518.             'lr0': (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
  519.             'lrf': (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
  520.             'momentum': (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
  521.             'weight_decay': (1, 0.0, 0.001),  # optimizer weight decay
  522.             'warmup_epochs': (1, 0.0, 5.0),  # warmup epochs (fractions ok)
  523.             'warmup_momentum': (1, 0.0, 0.95),  # warmup initial momentum
  524.             'warmup_bias_lr': (1, 0.0, 0.2),  # warmup initial bias lr
  525.             'box': (1, 0.02, 0.2),  # box loss gain
  526.             'cls': (1, 0.2, 4.0),  # cls loss gain
  527.             'cls_pw': (1, 0.5, 2.0),  # cls BCELoss positive_weight
  528.             'obj': (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
  529.             'obj_pw': (1, 0.5, 2.0),  # obj BCELoss positive_weight
  530.             'iou_t': (0, 0.1, 0.7),  # IoU training threshold
  531.             'anchor_t': (1, 2.0, 8.0),  # anchor-multiple threshold
  532.             'anchors': (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
  533.             'fl_gamma': (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
  534.             'hsv_h': (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
  535.             'hsv_s': (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
  536.             'hsv_v': (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
  537.             'degrees': (1, 0.0, 45.0),  # image rotation (+/- deg)
  538.             'translate': (1, 0.0, 0.9),  # image translation (+/- fraction)
  539.             'scale': (1, 0.0, 0.9),  # image scale (+/- gain)
  540.             'shear': (1, 0.0, 10.0),  # image shear (+/- deg)
  541.             'perspective': (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
  542.             'flipud': (1, 0.0, 1.0),  # image flip up-down (probability)
  543.             'fliplr': (0, 0.0, 1.0),  # image flip left-right (probability)
  544.             'mosaic': (1, 0.0, 1.0),  # image mixup (probability)
  545.             'mixup': (1, 0.0, 1.0),  # image mixup (probability)
  546.             'copy_paste': (1, 0.0, 1.0)}  # segment copy-paste (probability)
  548.         with open(opt.hyp, errors='ignore') as f:
  549.             hyp = yaml.safe_load(f)  # load hyps dict
  550.             if 'anchors' not in hyp:  # anchors commented in hyp.yaml
  551.                 hyp['anchors'] = 3
  552.         if opt.noautoanchor:
  553.             del hyp['anchors'], meta['anchors']
  554.         opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch
  555.         # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
  556.         evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
  557.         if opt.bucket:
  558.             # download evolve.csv if exists
  559.             subprocess.run([
  560.                 'gsutil',
  561.                 'cp',
  562.                 f'gs://{opt.bucket}/evolve.csv',
  563.                 str(evolve_csv),])
  565.         for _ in range(opt.evolve):  # generations to evolve
  566.             if evolve_csv.exists():  # if evolve.csv exists: select best hyps and mutate
  567.                 # Select parent(s)
  568.                 parent = 'single'  # parent selection method: 'single' or 'weighted'
  569.                 x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
  570.                 n = min(5, len(x))  # number of previous results to consider
  571.                 x = x[np.argsort(-fitness(x))][:n]  # top n mutations
  572.                 w = fitness(x) - fitness(x).min() + 1E-6  # weights (sum > 0)
  573.                 if parent == 'single' or len(x) == 1:
  574.                     # x = x[random.randint(0, n - 1)]  # random selection
  575.                     x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
  576.                 elif parent == 'weighted':
  577.                     x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
  579.                 # Mutate
  580.                 mp, s = 0.8, 0.2  # mutation probability, sigma
  581.                 npr = np.random
  582.                 npr.seed(int(time.time()))
  583.                 g = np.array([meta[k][0] for k in hyp.keys()])  # gains 0-1
  584.                 ng = len(meta)
  585.                 v = np.ones(ng)
  586.                 while all(v == 1):  # mutate until a change occurs (prevent duplicates)
  587.                     v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
  588.                 for i, k in enumerate(hyp.keys()):  # plt.hist(v.ravel(), 300)
  589.                     hyp[k] = float(x[i + 7] * v[i])  # mutate
  591.             # Constrain to limits
  592.             for k, v in meta.items():
  593.                 hyp[k] = max(hyp[k], v[1])  # lower limit
  594.                 hyp[k] = min(hyp[k], v[2])  # upper limit
  595.                 hyp[k] = round(hyp[k], 5)  # significant digits
  597.             # Train mutation
  598.             results = train(hyp.copy(), opt, device, callbacks)
  599.             callbacks = Callbacks()
  600.             # Write mutation results
  601.             keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss',
  602.                     'val/obj_loss', 'val/cls_loss')
  603.             print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
  605.         # Plot results
  606.         plot_evolve(evolve_csv)
  607.         LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
  608.                     f"Results saved to {colorstr('bold', save_dir)}\n"
  609.                     f'Usage example: $ python train.py --hyp {evolve_yaml}')
  612. def run(**kwargs):
  613.     # Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
  614.     opt = parse_opt(True)
  615.     for k, v in kwargs.items():
  616.         setattr(opt, k, v)
  617.     main(opt)
  618.     return opt
  621. if __name__ == '__main__':
  622.     opt = parse_opt()
  623.     main(opt)

（十）模型推断过程的程序编写

第一部分：编写 detect.py

  1. import argparse
  2. import os
  3. import platform
  4. import sys
  5. from pathlib import Path
  7. import torch
  9. FILE = Path(__file__).resolve()
  10. ROOT = FILE.parents[0]  # YOLOv5 root directory
  11. if str(ROOT) not in sys.path:
  12.     sys.path.append(str(ROOT))  # add ROOT to PATH
  13. ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
  15. from models.common import DetectMultiBackend
  16. from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
  17. from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
  18.                            increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
  19. from utils.plots import Annotator, colors, save_one_box
  20. from utils.torch_utils import select_device, smart_inference_mode
  23. @smart_inference_mode()
  24. def run(
  25.         weights=ROOT / r'D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\runs\train\exp15\weights\best.pt',
  26.         # model path or triton URL
  27.         # source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
  28.         source=ROOT / '0',  # file/dir/URL/glob/screen/0(webcam)
  29.         data=ROOT / 'car-parts/IP_parameter.yaml',  # dataset.yaml path
  30.         imgsz=(640, 640),  # inference size (height, width)
  31.         conf_thres=0.25,  # confidence threshold
  32.         iou_thres=0.45,  # NMS IOU threshold
  33.         max_det=1000,  # maximum detections per image
  34.         device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
  35.         view_img=False,  # show results
  36.         save_txt=False,  # save results to *.txt
  37.         save_conf=False,  # save confidences in --save-txt labels
  38.         save_crop=False,  # save cropped prediction boxes
  39.         nosave=False,  # do not save images/videos
  40.         classes=None,  # filter by class: --class 0, or --class 0 2 3
  41.         agnostic_nms=False,  # class-agnostic NMS
  42.         augment=False,  # augmented inference
  43.         visualize=False,  # visualize features
  44.         update=False,  # update all models
  45.         project=ROOT / 'runs/detect',  # save results to project/name
  46.         name='exp',  # save results to project/name
  47.         exist_ok=False,  # existing project/name ok, do not increment
  48.         line_thickness=3,  # bounding box thickness (pixels)
  49.         hide_labels=False,  # hide labels
  50.         hide_conf=False,  # hide confidences
  51.         half=False,  # use FP16 half-precision inference
  52.         dnn=False,  # use OpenCV DNN for ONNX inference
  53.         vid_stride=1,  # video frame-rate stride
  54. ):
  55.     source = str(source)
  56.     save_img = not nosave and not source.endswith('.txt')  # save inference images
  57.     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
  58.     is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
  59.     webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
  60.     screenshot = source.lower().startswith('screen')
  61.     if is_url and is_file:
  62.         source = check_file(source)  # download
  64.     # Directories
  65.     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
  66.     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
  68.     # Load model
  69.     device = select_device(device)
  70.     model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
  71.     stride, names, pt = model.stride, model.names, model.pt
  72.     imgsz = check_img_size(imgsz, s=stride)  # check image size
  74.     # Dataloader
  75.     bs = 1  # batch_size
  76.     if webcam:
  77.         view_img = check_imshow(warn=True)
  78.         dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
  79.         bs = len(dataset)
  80.     elif screenshot:
  81.         dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
  82.     else:
  83.         dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
  84.     vid_path, vid_writer = [None] * bs, [None] * bs
  86.     # Run inference
  87.     model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
  88.     seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
  89.     for path, im, im0s, vid_cap, s in dataset:
  90.         with dt[0]:
  91.             im = torch.from_numpy(im).to(model.device)
  92.             im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
  93.             im /= 255  # 0 - 255 to 0.0 - 1.0
  94.             if len(im.shape) == 3:
  95.                 im = im[None]  # expand for batch dim
  97.         # Inference
  98.         with dt[1]:
  99.             visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
  100.             pred = model(im, augment=augment, visualize=visualize)
  102.         # NMS
  103.         with dt[2]:
  104.             pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
  106.         # Second-stage classifier (optional)
  107.         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
  109.         # Process predictions
  110.         for i, det in enumerate(pred):  # per image
  111.             seen += 1
  112.             if webcam:  # batch_size >= 1
  113.                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
  114.                 s += f'{i}: '
  115.             else:
  116.                 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
  118.             p = Path(p)  # to Path
  119.             save_path = str(save_dir / p.name)  # im.jpg
  120.             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
  121.             s += '%gx%g ' % im.shape[2:]  # print string
  123.             #####################################################################################################
  125.             location_center_dir = str(save_dir) + '/location_center'
  126.             if not os.path.exists(location_center_dir):
  127.                 os.makedirs(location_center_dir)
  128.             location_center_path = location_center_dir + '\\' + str(p.stem) + (
  129.                 '' if dataset.mode == 'image' else f'_{frame}')  # location_center.txt
  130.             flocation = open(location_center_path + '.txt', 'w')  # 保存检测框中点
  132.             #####################################################################################################
  134.             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
  135.             imc = im0.copy() if save_crop else im0  # for save_crop
  136.             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
  137.             if len(det):
  138.                 # Rescale boxes from img_size to im0 size
  139.                 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
  141.                 # Print results
  142.                 for c in det[:, 5].unique():
  143.                     n = (det[:, 5] == c).sum()  # detections per class
  144.                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
  146.                 # Write results
  147.                 for *xyxy, conf, cls in reversed(det):
  149.                     #####################################################################################################
  151.                     x0 = (int(xyxy[0].item()) + int(xyxy[2].item())) / 2
  152.                     y0 = (int(xyxy[1].item()) + int(xyxy[3].item())) / 2  # 中心点坐标(x0, y0)
  153.                     class_index = cls  # 获取属性
  154.                     object_name = names[int(cls)]  # 获取标签名
  155.                     flocation.write(object_name + ': ' + str(x0) + ', ' + str(y0) + '\n')
  157.                     #####################################################################################################
  159.                     if save_txt:  # Write to file
  160.                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
  161.                         line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
  162.                         with open(f'{txt_path}.txt', 'a') as f:
  163.                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
  165.                     if save_img or save_crop or view_img:  # Add bbox to image
  166.                         c = int(cls)  # integer class
  167.                         label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
  168.                         annotator.box_label(xyxy, label, color=colors(c, True))
  169.                     if save_crop:
  170.                         save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
  171.             #####################################################################################################
  172.             flocation.close()
  173.             #####################################################################################################
  174.             # Stream results
  175.             im0 = annotator.result()
  176.             if view_img:
  177.                 if platform.system() == 'Linux' and p not in windows:
  178.                     windows.append(p)
  179.                     cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
  180.                     cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
  181.                 cv2.imshow(str(p), im0)
  182.                 cv2.waitKey(1)  # 1 millisecond
  184.             # Save results (image with detections)
  185.             if save_img:
  186.                 if dataset.mode == 'image':
  187.                     cv2.imwrite(save_path, im0)
  188.                 else:  # 'video' or 'stream'
  189.                     if vid_path[i] != save_path:  # new video
  190.                         vid_path[i] = save_path
  191.                         if isinstance(vid_writer[i], cv2.VideoWriter):
  192.                             vid_writer[i].release()  # release previous video writer
  193.                         if vid_cap:  # video
  194.                             fps = vid_cap.get(cv2.CAP_PROP_FPS)
  195.                             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  196.                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  197.                         else:  # stream
  198.                             fps, w, h = 30, im0.shape[1], im0.shape[0]
  199.                         save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
  200.                         vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
  201.                     vid_writer[i].write(im0)
  203.         # Print time (inference-only)
  204.         LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
  206.     # Print results
  207.     t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
  208.     LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
  209.     if save_txt or save_img:
  210.         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  211.         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
  212.     if update:
  213.         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
  216. def parse_opt():
  217.     parser = argparse.ArgumentParser()
  218.     # parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL')
  219.     parser.add_argument('--weights', nargs='+', type=str,
  220.                         default=ROOT / r'D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\runs\train\exp3\weights\best.pt',
  221.                         help='model path or triton URL')
  222.     # parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
  223.     parser.add_argument('--source', type=str,
  224.                         default=ROOT / r'D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\V2\datasets\images\train',
  225.                         help='file/dir/URL/glob/screen/0(webcam)')
  226.     parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
  227.     parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
  228.     parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
  229.     parser.add_argument('--iou-thres', type=float, default=0, help='NMS IoU threshold')
  230.     parser.add_argument('--max-det', type=int, default=10, help='maximum detections per image')
  231.     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  232.     parser.add_argument('--view-img', action='store_true', help='show results')
  233.     parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  234.     parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  235.     parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
  236.     parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
  237.     parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
  238.     parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
  239.     parser.add_argument('--augment', action='store_true', help='augmented inference')
  240.     parser.add_argument('--visualize', action='store_true', help='visualize features')
  241.     parser.add_argument('--update', action='store_true', help='update all models')
  242.     parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
  243.     parser.add_argument('--name', default='exp', help='save results to project/name')
  244.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  245.     parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
  246.     parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
  247.     parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
  248.     parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
  249.     parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
  250.     parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
  251.     opt = parser.parse_args()
  252.     opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
  253.     print_args(vars(opt))
  254.     return opt
  257. def main(opt):
  258.     check_requirements(exclude=('tensorboard', 'thop'))
  259.     run(**vars(opt))
  262. if __name__ == '__main__':
  263.     opt = parse_opt()
  264.     main(opt)

第二部分：在无人车终端运行代码，调用物理摄像头，查看具体效果

（十一）使用onnx及TensorRT将模型进行推理加速

第一步：将pt模型转化为onnx模型

① 安装onnx

  1. pip install onnx
  1. pip install onnxruntime

② 使用 export.py 导出模型为ONNX

  1. python export.py --weights D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\runs\train\exp3\weights\best.pt --img-size 640 --batch-size 1 --include onnx

第二步：TensorRT环境安装及配置

① 下载对应TensorRT版本

  1. https://developer.nvidia.com/nvidia-tensorrt-8x-download

② 解压 TensorRT

③ 配置环境变量

将TensorRT解压位置\lib 加入系统环境变量

将TensorRT解压位置\lib下的dll文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin目录下

④ 测试示例代码

用VS2019打开sampleOnnxMNIST示例(D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\samples\sampleOnnxMNIST）

将D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\lib加入 VC++目录–>可执行文件目录

将D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\include加入C/C++ --> 常规 --> 附加包含目录

将D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\lib加入 VC++目录–>库目录

将nvinfer.lib、nvinfer_plugin.lib、nvonnxparser.lib和nvparsers.lib加入链接器–>输入–>附加依赖项

⑤ 安装 pycuda

https://www.lfd.uci.edu/~gohlke/pythonlibs/?cm_mc_uid=08085305845514542921829&cm_mc_sid_50200000=1456395916&cm_mc_uid=08085305845514542921829&cm_mc_sid_50200000=1456395916#pycuda

  1. pip install "D:\pycuda\pycuda-2022.1+cuda116-cp38-cp38-win_amd64.whl"

⑥ python 环境配置 TensorRT

  1. pip install "D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\python\tensorrt-8.6.1-cp38-none-win_amd64.whl"

⑦ python 环境测试

  1. python "D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\samples\python\network_api_pytorch_mnist\sample.py"

第三步：使用 TensorRT 编译 onnx 文件，转换成 .trt 后缀文件

D:\TensorRT\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\bin\trtexec.exe --onnx=D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\runs\train\exp3\weights\best.onnx --saveEngine=D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\runs\train\exp3\weights\best.trt --buildOnly

第四步：查看训练图片的维度

check_picture.py

  1. import cv2
  3. # 读取图像
  4. image_path = r'D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\V2\datasets\images\train\2023-09-18-160239.jpg '
  5. image = cv2.imread(image_path)
  7. # 检查图像形状
  8. image_shape = image.shape
  10. # 打印图像形状
  11. print("图像形状:", image_shape)
  13. # 如果您只关心图像的高度，宽度和通道数，可以使用以下方式获取
  14. height, width, channels = image_shape
  15. print("高度:", height)
  16. print("宽度:", width)
  17. print("通道数:", channels)

第五步：编写加速推理脚本

compare.py

  1. import numpy as np
  2. import pycuda.driver as cuda
  3. import pycuda.autoinit
  4. import tensorrt as trt
  5. import cv2
  6. import time
  8. # 加载.trt文件
  9. trt_file_path = r'D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\runs\train\exp3\weights\best.trt'
  10. with open(trt_file_path, 'rb') as f, trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime:
  11.     engine = runtime.deserialize_cuda_engine(f.read())
  13. # 创建执行上下文
  14. context = engine.create_execution_context()
  16. # 分配输入和输出内存
  17. input_size = trt.volume(context.get_binding_shape(0))
  18. output_size = trt.volume(context.get_binding_shape(1))
  20. # 在GPU上分配内存
  21. input_host = cuda.pagelocked_empty(input_size, dtype=np.float32)
  22. output_host = cuda.pagelocked_empty(output_size, dtype=np.float32)
  23. input_device = cuda.mem_alloc(input_host.nbytes)
  24. output_device = cuda.mem_alloc(output_host.nbytes)
  26. # 准备输入数据
  27. # 读取彩色图像
  28. image_path = r'D:\BaiduNetdiskWorkspace\licenses\machine_learning\yolov5\yolov5-master\V2\datasets\images\train\2023-09-18-160239.jpg '
  29. image = cv2.imread(image_path)
  31. # 调整数据形状以匹配模型期望的输入形状
  32. input_data = image.astype(np.float32) / 255.0  # 归一化（假设模型期望的输入范围是 [0, 1]）
  34. # 使用cv2.resize调整图像大小
  35. resized_image = cv2.resize(input_data, (640, 640))
  37. input_data = np.transpose(resized_image, (2, 0, 1))  # 将通道移到正确的位置
  38. input_data = np.expand_dims(input_data, axis=0)  # 添加批处理维度
  40. # 确保输入数据的长度与模型期望的输入大小一致
  41. if input_data.size != input_size:
  42.     raise ValueError(f"Input data size ({input_data.size}) does not match the expected input size ({input_size})")
  44. np.copyto(input_host, input_data.ravel())
  45. cuda.memcpy_htod(input_device, input_host)
  47. # 计时开始
  48. start_time = time.time()
  50. # 执行推理
  51. context.execute_v2(bindings=[int(input_device), int(output_device)])
  53. # 计时结束
  54. end_time = time.time()
  56. # 获取输出
  57. cuda.memcpy_dtoh(output_host, output_device)
  59. # 处理输出数据
  60. result = output_host.reshape(context.get_binding_shape(1))
  62. # 打印结果
  63. print(result)
  65. # 打印推理时间
  66. inference_time = end_time - start_time
  67. print(f"Inference time: {inference_time} seconds")

第六步：对比测试

加速前

加速后

（十二）问题汇总

（1）error MSB8036：找不到 Windows SDK 版本......

参考链接

error MSB8036：找不到 Windows SDK 版本···_microsoft.cpp.windowssdk.targe ts(46,5): error msb-CSDN博客

（2）too many values to unpack

参考链接

too many values to unpack (expected 4)_UC_Gundam的博客-CSDN博客

（3）运行示例时，提示找不到 MNIST数据

参考链接

TensorRT之安装与测试（Windows和Linux环境下安装TensorRT）_判断tensorrt是否可以正常使用-CSDN博客

（4）加载 libnvinfer.so.7报错

参考链接

TensorRT之安装与测试（Windows和Linux环境下安装TensorRT）_判断tensorrt是否可以正常使用-CSDN博客

原文链接：https://www.cnblogs.com/zylyehuo/p/17832064.html

 友情链接：直通硅谷 　点职佳 　北美留学生论坛

人工智能基础热门文章

  • yolov5 AssertionError: Image Not Found解决方案daily-
  • 如何注册chatgpt，如何使用chatgpt，以及chatgpt无法访问的原因。chatgpt问题总结。
  • 并行Louvain社区检测算法 - orion-orion
  • 《Deep Learning》(深度学习)中文版PDF免费下载
  • 斗地主AI出牌助手--在线调用斗地主AI，实现自动斗地主
  • ChatGPT 1.0.0安卓分析，仅限国内分享
  • ChatGPT3.5使用体验
  • 0902-用GAN生成动漫头像
  • 记录：tf.saved_model 模块的简单使用（TensorFlow 模型存储与恢复）
  • ML.NET教程之客户细分(聚类问题)

    • 人工智能基础推荐文章

  • 【LLM训练系列】NanoGPT源码详解和中文GPT训练实践
  • 解密Prompt系列36. Prompt结构化编写和最优化算法UNIPROMPT
  • RealNet：从数据生成到模型训练的最新工业异常检测 | CVPR 2024
  • 再探GraphRAG：如何提升LLM总结能力？
  • AI 大模型时代呼唤新一代基础设施，DataOps 2.0和调度编排愈发重要
  • QWen2-72B-Instruct模型安装部署过程
  • AI编程助手那些事儿
  • FP分数规划在无线通信中的应用
  • ambari+ bigtop 编译、打包、部署步骤总览
  • ComfyUI插件：ComfyUI Impact 节点（二）

    • 本站QQ群：前端 618073944 | Java 606181507 | Python 626812652 | C/C++ 612253063 | 微信 634508462 | 苹果 692586424 | C#/.net 182808419 | PHP 305140648 | 运维 608723728

    W3xue 的所有内容仅供测试，对任何法律问题及风险不承担任何责任。通过使用本站内容随之而来的风险与本站无关。
    关于我们  |  意见建议  |  捐助我们  |  报错有奖  |  广告合作、友情链接（目前9元/月）请联系QQ：27243702 沸活量
    皖ICP备17017327号-2　皖公网安备34020702000426号