devbox
很楠不爱3
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

基于YOLOv5、YOLOv8的火灾检测（超实用项目）_国内场景 电动车失火yolov5检测

作者：很楠不爱3 | 2024-06-14 14:54:30

赞

踩

国内场景 电动车失火yolov5检测

目录

1.简介

2.YOLO算法

3.基于YOLOv5、YOLOv8的火灾检测

 

 视频已上传b站

YOLOv5/YOLOv8的火灾检测（超实用项目）_哔哩哔哩_bilibili

本文为系列专栏，包括各种YOLO检测算法项目、追踪算法项目、双目视觉、深度结构光相机测距测速三维测量项目等

专栏持续更新中，有需要的小伙伴可私聊，接项目定制。

1.简介

随着科技的不断发展，人工智能技术在各个领域得到广泛应用。其中，计算机视觉是人工智能领域中的一个重要分支，它主要研究如何使机器“看”和“理解”图像或视频。在计算机视觉领域，目标检测是一个关键问题，它涉及识别图像或视频中的特定对象，并确定它们的位置。火灾检测作为目标检测的一个重要应用领域，对于及时发现火灾、减少人员伤亡和财产损失具有重要意义。

本项目旨在基于YOLOv5和YOLOv8这两个先进的目标检测模型，开展火灾检测的研究和应用。YOLO（You Only Look Once）是一种实时目标检测算法，它将目标检测任务转化为一个回归问题，通过单次前向传递神经网络即可得到图像中所有目标的类别和位置。YOLOv5是YOLO系列中的最新版本，它在精度和速度之间取得了很好的平衡，被广泛应用于各种实时目标检测任务。

在本项目中，我们将探讨火灾检测领域的挑战和需求，介绍YOLOv5和YOLOv8的基本原理和算法结构，以及在实际火灾检测场景中的应用。通过本项目的研究，我们希望能够为提高火灾检测的准确性和效率，保障人们的生命财产安全，做出积极贡献。

希望本项目能够为火灾检测领域的研究和实际应用提供有益的参考和启示，推动人工智能技术在火灾安全领域的进一步发展和应用。

2.YOLO算法

YOLO（You Only Look Once）是一种高效的实时目标检测算法，它将目标检测任务转化为一个回归问题，通过单次前向传递神经网络即可得到图像中所有目标的类别和位置。相较于传统的目标检测方法，YOLO具有更快的速度和较高的准确性，使其成为计算机视觉领域中的重要算法之一。

YOLO算法的基本思想是将输入图像划分为一个固定大小的网格（grid），每个网格负责预测图像中是否包含目标以及目标的位置和类别。与传统的滑动窗口方法不同，YOLO将目标检测任务转化为一个回归问题，同时预测所有目标的位置和类别，避免了重复计算，因此速度更快。

 

 

以下是YOLO算法的主要特点和步骤：

  1. 划分网格： 将输入图像划分为SxS个网格，每个网格负责预测该网格内是否包含目标。

  2. 预测框和类别： 每个网格预测B个边界框（bounding boxes）以及每个边界框的置信度（confidence）和类别概率。置信度表示边界框的准确性，类别概率表示目标属于不同类别的概率。

  3. 计算损失函数： YOLO使用多任务损失函数，包括边界框坐标的回归损失、置信度的损失（包括目标是否存在的损失和目标位置的精度损失）、类别概率的损失。通过最小化这些损失，网络可以学习到准确的目标位置和类别信息。

  4. 非极大值抑制（NMS）： 在预测结果中，可能存在多个边界框对同一个目标的重复检测。为了去除这些重叠的边界框，使用NMS算法来选择具有最高置信度的边界框，并消除与其IoU（交并比）高于阈值的其他边界框。

  5. 输出结果： 最终，YOLO输出图像中所有目标的位置和类别信息，以及它们的置信度分数。

YOLO的优势在于它的速度和准确性，它能够实时处理高分辨率的图像，并且在不同尺度和大小的目标上具有很好的泛化能力。这使得YOLO广泛应用于实时目标检测、视频分析、自动驾驶等领域。

3.基于YOLOv5、YOLOv8的火灾检测

部分代码展示

gui界面主代码

  1. from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMenu, QAction
  2. from main_win.win import Ui_mainWindow
  3. from PyQt5.QtCore import Qt, QPoint, QTimer, QThread, pyqtSignal
  4. from PyQt5.QtGui import QImage, QPixmap, QPainter, QIcon
  5.  
  6. import sys
  7. import os
  8. import json
  9. import numpy as np
  10. import torch
  11. import torch.backends.cudnn as cudnn
  12. import os
  13. import time
  14. import cv2
  15.  
  16. from models.experimental import attempt_load
  17. from utils.datasets import LoadImages, LoadWebcam
  18. from utils.CustomMessageBox import MessageBox
  19. # LoadWebcam 的最后一个返回值改为 self.cap
  20. from utils.general import check_img_size, check_requirements, check_imshow, colorstr, non_max_suppression, \
  21.     apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
  22. from utils.plots import colors, plot_one_box, plot_one_box_PIL
  23. from utils.torch_utils import select_device, load_classifier, time_sync
  24. from utils.capnums import Camera
  25. from dialog.rtsp_win import Window
  26.  
  27.  
  28. class DetThread(QThread):
  29.     send_img = pyqtSignal(np.ndarray)
  30.     send_raw = pyqtSignal(np.ndarray)
  31.     send_statistic = pyqtSignal(dict)
  32.     # 发送信号：正在检测/暂停/停止/检测结束/错误报告
  33.     send_msg = pyqtSignal(str)
  34.     send_percent = pyqtSignal(int)
  35.     send_fps = pyqtSignal(str)
  36.  
  37.     def __init__(self):
  38.         super(DetThread, self).__init__()
  39.         self.weights = './yolov5s.pt'           # 设置权重
  40.         self.current_weight = './yolov5s.pt'    # 当前权重
  41.         self.source = '0'                       # 视频源
  42.         self.conf_thres = 0.25                  # 置信度
  43.         self.iou_thres = 0.45                   # iou
  44.         self.jump_out = False                   # 跳出循环
  45.         self.is_continue = True                 # 继续/暂停
  46.         self.percent_length = 1000              # 进度条
  47.         self.rate_check = True                  # 是否启用延时
  48.         self.rate = 100                         # 延时HZ
  49.         self.save_fold = './result'             # 保存文件夹
  50.  
  51.     @torch.no_grad()
  52.     def run(self,
  53.             imgsz=640,  # inference size (pixels)
  54.             max_det=1000,  # maximum detections per image
  55.             device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
  56.             view_img=True,  # show results
  57.             save_txt=False,  # save results to *.txt
  58.             save_conf=False,  # save confidences in --save-txt labels
  59.             save_crop=False,  # save cropped prediction boxes
  60.             nosave=False,  # do not save images/videos
  61.             classes=None,  # filter by class: --class 0, or --class 0 2 3
  62.             agnostic_nms=False,  # class-agnostic NMS
  63.             augment=False,  # augmented inference
  64.             visualize=False,  # visualize features
  65.             update=False,  # update all models
  66.             project='runs/detect',  # save results to project/name
  67.             name='exp',  # save results to project/name
  68.             exist_ok=False,  # existing project/name ok, do not increment
  69.             line_thickness=3,  # bounding box thickness (pixels)
  70.             hide_labels=False,  # hide labels
  71.             hide_conf=False,  # hide confidences
  72.             half=False,  # use FP16 half-precision inference
  73.             ):
  74.  
  75.         # Initialize
  76.         try:
  77.             device = select_device(device)
  78.             half &= device.type != 'cpu'  # half precision only supported on CUDA
  79.  
  80.             # Load model
  81.             model = attempt_load(self.weights, map_location=device)  # load FP32 model
  82.             num_params = 0
  83.             for param in model.parameters():
  84.                 num_params += param.numel()
  85.             stride = int(model.stride.max())  # model stride
  86.             imgsz = check_img_size(imgsz, s=stride)  # check image size
  87.             names = model.module.names if hasattr(model, 'module') else model.names  # get class names
  88.             if half:
  89.                 model.half()  # to FP16
  90.  
  91.             # Dataloader
  92.             if self.source.isnumeric() or self.source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')):
  93.                 view_img = check_imshow()
  94.                 cudnn.benchmark = True  # set True to speed up constant image size inference
  95.                 dataset = LoadWebcam(self.source, img_size=imgsz, stride=stride)
  96.                 # bs = len(dataset)  # batch_size
  97.             else:
  98.                 dataset = LoadImages(self.source, img_size=imgsz, stride=stride)
  99.  
  100.             # Run inference
  101.             if device.type != 'cpu':
  102.                 model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
  103.             count = 0
  104.             # 跳帧检测
  105.             jump_count = 0
  106.             start_time = time.time()
  107.             dataset = iter(dataset)
  108.  
  109.             while True:
  110.                 # 手动停止
  111.                 if self.jump_out:
  112.                     self.vid_cap.release()
  113.                     self.send_percent.emit(0)
  114.                     self.send_msg.emit('停止')
  115.                     if hasattr(self, 'out'):
  116.                         self.out.release()
  117.                     break
  118.                 # 临时更换模型
  119.                 if self.current_weight != self.weights:
  120.                     # Load model
  121.                     model = attempt_load(self.weights, map_location=device)  # load FP32 model
  122.                     num_params = 0
  123.                     for param in model.parameters():
  124.                         num_params += param.numel()
  125.                     stride = int(model.stride.max())  # model stride
  126.                     imgsz = check_img_size(imgsz, s=stride)  # check image size
  127.                     names = model.module.names if hasattr(model, 'module') else model.names  # get class names
  128.                     if half:
  129.                         model.half()  # to FP16
  130.                     # Run inference
  131.                     if device.type != 'cpu':
  132.                         model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
  133.                     self.current_weight = self.weights
  134.                 # 暂停开关
  135.                 if self.is_continue:
  136.                     path, img, im0s, self.vid_cap = next(dataset)
  137.                     # jump_count += 1
  138.                     # if jump_count % 5 != 0:
  139.                     #     continue
  140.                     count += 1
  141.                     # 每三十帧刷新一次输出帧率
  142.                     if count % 30 == 0 and count >= 30:
  143.                         fps = int(30/(time.time()-start_time))
  144.                         self.send_fps.emit('fps：'+str(fps))
  145.                         start_time = time.time()
  146.                     if self.vid_cap:
  147.                         percent = int(count/self.vid_cap.get(cv2.CAP_PROP_FRAME_COUNT)*self.percent_length)
  148.                         self.send_percent.emit(percent)
  149.                     else:
  150.                         percent = self.percent_length
  151.  
  152.                     statistic_dic = {name: 0 for name in names}
  153.                     img = torch.from_numpy(img).to(device)
  154.                     img = img.half() if half else img.float()  # uint8 to fp16/32
  155.                     img /= 255.0  # 0 - 255 to 0.0 - 1.0
  156.                     if img.ndimension() == 3:
  157.                         img = img.unsqueeze(0)
  158.  
  159.                     pred = model(img, augment=augment)[0]
  160.  
  161.                     # Apply NMS
  162.                     pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes, agnostic_nms, max_det=max_det)
  163.                     # Process detections
  164.                     for i, det in enumerate(pred):  # detections per image
  165.                         im0 = im0s.copy()
  166.  
  167.                         if len(det):
  168.                             # Rescale boxes from img_size to im0 size
  169.                             det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
  170.  
  171.                             # Write results
  172.                             for *xyxy, conf, cls in reversed(det):
  173.                                 c = int(cls)  # integer class
  174.                                 statistic_dic[names[c]] += 1
  175.                                 label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
  176.                                 # im0 = plot_one_box_PIL(xyxy, im0, label=label, color=colors(c, True), line_thickness=line_thickness)  # 中文标签画框，但是耗时会增加
  177.                                 plot_one_box(xyxy, im0, label=label, color=colors(c, True),
  178.                                              line_thickness=line_thickness)
  179.  
  180.                     # 控制视频发送频率
  181.                     if self.rate_check:
  182.                         time.sleep(1/self.rate)
  183.                     self.send_img.emit(im0)
  184.                     self.send_raw.emit(im0s if isinstance(im0s, np.ndarray) else im0s[0])
  185.                     self.send_statistic.emit(statistic_dic)
  186.                     # 如果自动录制
  187.                     if self.save_fold:
  188.                         os.makedirs(self.save_fold, exist_ok=True)  # 路径不存在，自动保存
  189.                         # 如果输入是图片
  190.                         if self.vid_cap is None:
  191.                             save_path = os.path.join(self.save_fold,
  192.                                                      time.strftime('%Y_%m_%d_%H_%M_%S',
  193.                                                                    time.localtime()) + '.jpg')
  194.                             cv2.imwrite(save_path, im0)
  195.                         else:
  196.                             if count == 1:  # 第一帧时初始化录制
  197.                                 # 以视频原始帧率进行录制
  198.                                 ori_fps = int(self.vid_cap.get(cv2.CAP_PROP_FPS))
  199.                                 if ori_fps == 0:
  200.                                     ori_fps = 25
  201.                                 # width = int(self.vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  202.                                 # height = int(self.vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  203.                                 width, height = im0.shape[1], im0.shape[0]
  204.                                 save_path = os.path.join(self.save_fold, time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime()) + '.mp4')
  205.                                 self.out = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), ori_fps,
  206.                                                            (width, height))
  207.                             self.out.write(im0)
  208.                     if percent == self.percent_length:
  209.                         print(count)
  210.                         self.send_percent.emit(0)
  211.                         self.send_msg.emit('检测结束')
  212.                         if hasattr(self, 'out'):
  213.                             self.out.release()
  214.                         # 正常跳出循环
  215.                         break
  216.  
  217.         except Exception as e:
  218.             self.send_msg.emit('%s' % e)

 YOLOv5主代码

  1. """Train a YOLOv5 model on a custom dataset
  3. Usage:
  4.     $ python path/to/train.py --data coco128.yaml --weights yolov5s.pt --img 640
  5. """
  6.  
  7. import argparse
  8. import logging
  9. import os
  10. import random
  11. import sys
  12. import time
  13. import warnings
  14. from copy import deepcopy
  15. from pathlib import Path
  16. from threading import Thread
  17.  
  18. import math
  19. import numpy as np
  20. import torch.distributed as dist
  21. import torch.nn as nn
  22. import torch.nn.functional as F
  23. import torch.optim as optim
  24. import torch.optim.lr_scheduler as lr_scheduler
  25. import torch.utils.data
  26. import yaml
  27. from torch.cuda import amp
  28. from torch.nn.parallel import DistributedDataParallel as DDP
  29. from torch.utils.tensorboard import SummaryWriter
  30. from tqdm import tqdm
  31.  
  32. FILE = Path(__file__).absolute()
  33. sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path
  34.  
  35. import val  # for end-of-epoch mAP
  36. from models.experimental import attempt_load
  37. from models.yolo import Model
  38. from utils.autoanchor import check_anchors
  39. from utils.datasets import create_dataloader
  40. from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
  41.     strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \
  42.     check_requirements, print_mutation, set_logging, one_cycle, colorstr
  43. from utils.google_utils import attempt_download
  44. from utils.loss import ComputeLoss
  45. from utils.plots import plot_images, plot_labels, plot_results, plot_evolution
  46. from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, de_parallel
  47. from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
  48. from utils.metrics import fitness
  49.  
  50. LOGGER = logging.getLogger(__name__)
  51. LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
  52. RANK = int(os.getenv('RANK', -1))
  53. WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
  54.  
  55.  
  56. def train(hyp,  # path/to/hyp.yaml or hyp dictionary
  57.           opt,
  58.           device,
  59.           ):
  60.     save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, = \
  61.         opt.save_dir, opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
  62.         opt.resume, opt.noval, opt.nosave, opt.workers
  63.  
  64.     # Directories
  65.     save_dir = Path(save_dir)
  66.     wdir = save_dir / 'weights'
  67.     wdir.mkdir(parents=True, exist_ok=True)  # make dir
  68.     last = wdir / 'last.pt'
  69.     best = wdir / 'best.pt'
  70.     results_file = save_dir / 'results.txt'
  71.  
  72.     # Hyperparameters
  73.     if isinstance(hyp, str):
  74.         with open(hyp) as f:
  75.             hyp = yaml.safe_load(f)  # load hyps dict
  76.     LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
  77.  
  78.     # Save run settings
  79.     with open(save_dir / 'hyp.yaml', 'w') as f:
  80.         yaml.safe_dump(hyp, f, sort_keys=False)
  81.     with open(save_dir / 'opt.yaml', 'w') as f:
  82.         yaml.safe_dump(vars(opt), f, sort_keys=False)
  83.  
  84.     # Configure
  85.     plots = not evolve  # create plots
  86.     cuda = device.type != 'cpu'
  87.     init_seeds(1 + RANK)
  88.     with open(data) as f:
  89.         data_dict = yaml.safe_load(f)  # data dict
  90.  
  91.     # Loggers
  92.     loggers = {'wandb': None, 'tb': None}  # loggers dict
  93.     if RANK in [-1, 0]:
  94.         # TensorBoard
  95.         if not evolve:
  96.             prefix = colorstr('tensorboard: ')
  97.             LOGGER.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/")
  98.             loggers['tb'] = SummaryWriter(str(save_dir))
  99.  
  100.         # W&B
  101.         opt.hyp = hyp  # add hyperparameters
  102.         run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
  103.         run_id = run_id if opt.resume else None  # start fresh run if transfer learning
  104.         wandb_logger = WandbLogger(opt, save_dir.stem, run_id, data_dict)
  105.         loggers['wandb'] = wandb_logger.wandb
  106.         if loggers['wandb']:
  107.             data_dict = wandb_logger.data_dict
  108.             weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp  # may update weights, epochs if resuming
  109.  
  110.     nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
  111.     names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
  112.     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, data)  # check
  113.     is_coco = data.endswith('coco.yaml') and nc == 80  # COCO dataset
  114.  
  115.     # Model
  116.     pretrained = weights.endswith('.pt')
  117.     if pretrained:
  118.         with torch_distributed_zero_first(RANK):
  119.             weights = attempt_download(weights)  # download if not found locally
  120.         ckpt = torch.load(weights, map_location=device)  # load checkpoint
  121.         model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  122.         exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []  # exclude keys
  123.         state_dict = ckpt['model'].float().state_dict()  # to FP32
  124.         state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
  125.         model.load_state_dict(state_dict, strict=False)  # load
  126.         LOGGER.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
  127.     else:
  128.         model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  129.     with torch_distributed_zero_first(RANK):
  130.         check_dataset(data_dict)  # check
  131.     train_path = data_dict['train']
  132.     val_path = data_dict['val']
  133.  
  134.     # Freeze
  135.     freeze = []  # parameter names to freeze (full or partial)
  136.     for k, v in model.named_parameters():
  137.         v.requires_grad = True  # train all layers
  138.         if any(x in k for x in freeze):
  139.             print('freezing %s' % k)
  140.             v.requires_grad = False
  141.  
  142.     # Optimizer
  143.     nbs = 64  # nominal batch size
  144.     accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
  145.     hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
  146.     LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
  147.  
  148.     pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
  149.     for k, v in model.named_modules():
  150.         if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
  151.             pg2.append(v.bias)  # biases
  152.         if isinstance(v, nn.BatchNorm2d):
  153.             pg0.append(v.weight)  # no decay
  154.         elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
  155.             pg1.append(v.weight)  # apply decay
  156.  
  157.     if opt.adam:
  158.         optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum
  159.     else:
  160.         optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
  161.  
  162.     optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})  # add pg1 with weight_decay
  163.     optimizer.add_param_group({'params': pg2})  # add pg2 (biases)
  164.     LOGGER.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
  165.     del pg0, pg1, pg2
  166.  
  167.     # Scheduler https://arxiv.org/pdf/1812.01187.pdf
  168.     # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
  169.     if opt.linear_lr:
  170.         lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
  171.     else:
  172.         lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
  173.     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
  174.     # plot_lr_scheduler(optimizer, scheduler, epochs)
  175.  
  176.     # EMA
  177.     ema = ModelEMA(model) if RANK in [-1, 0] else None
  178.  
  179.     # Resume
  180.     start_epoch, best_fitness = 0, 0.0
  181.     if pretrained:
  182.         # Optimizer
  183.         if ckpt['optimizer'] is not None:
  184.             optimizer.load_state_dict(ckpt['optimizer'])
  185.             best_fitness = ckpt['best_fitness']
  186.  
  187.         # EMA
  188.         if ema and ckpt.get('ema'):
  189.             ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
  190.             ema.updates = ckpt['updates']
  191.  
  192.         # Results
  193.         if ckpt.get('training_results') is not None:
  194.             results_file.write_text(ckpt['training_results'])  # write results.txt
  195.  
  196.         # Epochs
  197.         start_epoch = ckpt['epoch'] + 1
  198.         if resume:
  199.             assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
  200.         if epochs < start_epoch:
  201.             LOGGER.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
  202.                         (weights, ckpt['epoch'], epochs))
  203.             epochs += ckpt['epoch']  # finetune additional epochs
  204.  
  205.         del ckpt, state_dict
  206.  
  207.     # Image sizes
  208.     gs = max(int(model.stride.max()), 32)  # grid size (max stride)
  209.     nl = model.model[-1].nl  # number of detection layers (used for scaling hyp['obj'])
  210.     imgsz = check_img_size(opt.imgsz, gs)  # verify imgsz is gs-multiple
  211.  
  212.     # DP mode
  213.     if cuda and RANK == -1 and torch.cuda.device_count() > 1:
  214.         logging.warning('DP not recommended, instead use torch.distributed.run for best DDP Multi-GPU results.\n'
  215.                         'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
  216.         model = torch.nn.DataParallel(model)
  217.  
  218.     # SyncBatchNorm
  219.     if opt.sync_bn and cuda and RANK != -1:
  220.         raise Exception('can not train with --sync-bn, known issue https://github.com/ultralytics/yolov5/issues/3998')
  221.         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
  222.         LOGGER.info('Using SyncBatchNorm()')
  223.  
  224.     # Trainloader
  225.     train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
  226.                                               hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=RANK,
  227.                                               workers=workers, image_weights=opt.image_weights, quad=opt.quad,
  228.                                               prefix=colorstr('train: '))
  229.     mlc = np.concatenate(dataset.labels, 0)[:, 0].max()  # max label class
  230.     nb = len(train_loader)  # number of batches
  231.     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, data, nc - 1)
  232.  
  233.     # Process 0
  234.     if RANK in [-1, 0]:
  235.         val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls,
  236.                                        hyp=hyp, cache=opt.cache_images and not noval, rect=True, rank=-1,
  237.                                        workers=workers, pad=0.5,
  238.                                        prefix=colorstr('val: '))[0]
  239.  
  240.         if not resume:
  241.             labels = np.concatenate(dataset.labels, 0)
  242.             # c = torch.tensor(labels[:, 0])  # classes
  243.             # cf = torch.bincount(c.long(), minlength=nc) + 1.  # frequency
  244.             # model._initialize_biases(cf.to(device))
  245.             if plots:
  246.                 plot_labels(labels, names, save_dir, loggers)
  247.  
  248.             # Anchors
  249.             if not opt.noautoanchor:
  250.                 check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
  251.             model.half().float()  # pre-reduce anchor precision
  252.  
  253.     # DDP mode
  254.     if cuda and RANK != -1:
  255.         model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
  256.  
  257.     # Model parameters
  258.     hyp['box'] *= 3. / nl  # scale to layers
  259.     hyp['cls'] *= nc / 80. * 3. / nl  # scale to classes and layers
  260.     hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl  # scale to image size and layers
  261.     hyp['label_smoothing'] = opt.label_smoothing
  262.     model.nc = nc  # attach number of classes to model
  263.     model.hyp = hyp  # attach hyperparameters to model
  264.     model.gr = 1.0  # iou loss ratio (obj_loss = 1.0 or iou)
  265.     model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
  266.     model.names = names
  267.  
  268.     # Start training
  269.     t0 = time.time()
  270.     nw = max(round(hyp['warmup_epochs'] * nb), 1000)  # number of warmup iterations, max(3 epochs, 1k iterations)
  271.     # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
  272.     last_opt_step = -1
  273.     maps = np.zeros(nc)  # mAP per class
  274.     results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
  275.     scheduler.last_epoch = start_epoch - 1  # do not move
  276.     scaler = amp.GradScaler(enabled=cuda)
  277.     compute_loss = ComputeLoss(model)  # init loss class
  278.     LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
  279.                 f'Using {train_loader.num_workers} dataloader workers\n'
  280.                 f'Logging results to {save_dir}\n'
  281.                 f'Starting training for {epochs} epochs...')
  282.     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
  283.         model.train()
  284.  
  285.         # Update image weights (optional)
  286.         if opt.image_weights:
  287.             # Generate indices
  288.             if RANK in [-1, 0]:
  289.                 cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights
  290.                 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
  291.                 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
  292.             # Broadcast if DDP
  293.             if RANK != -1:
  294.                 indices = (torch.tensor(dataset.indices) if RANK == 0 else torch.zeros(dataset.n)).int()
  295.                 dist.broadcast(indices, 0)
  296.                 if RANK != 0:
  297.                     dataset.indices = indices.cpu().numpy()
  298.  
  299.         # Update mosaic border
  300.         # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
  301.         # dataset.mosaic_border = [b - imgsz, -b]  # height, width borders
  302.  
  303.         mloss = torch.zeros(4, device=device)  # mean losses
  304.         if RANK != -1:
  305.             train_loader.sampler.set_epoch(epoch)
  306.         pbar = enumerate(train_loader)
  307.         LOGGER.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size'))
  308.         if RANK in [-1, 0]:
  309.             pbar = tqdm(pbar, total=nb)  # progress bar
  310.         optimizer.zero_grad()
  311.         for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
  312.             ni = i + nb * epoch  # number integrated batches (since train start)
  313.             imgs = imgs.to(device, non_blocking=True).float() / 255.0  # uint8 to float32, 0-255 to 0.0-1.0
  314.  
  315.             # Warmup
  316.             if ni <= nw:
  317.                 xi = [0, nw]  # x interp
  318.                 # model.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
  319.                 accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
  320.                 for j, x in enumerate(optimizer.param_groups):
  321.                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
  322.                     x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
  323.                     if 'momentum' in x:
  324.                         x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
  325.  
  326.             # Multi-scale
  327.             if opt.multi_scale:
  328.                 sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size
  329.                 sf = sz / max(imgs.shape[2:])  # scale factor
  330.                 if sf != 1:
  331.                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
  332.                     imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
  333.  
  334.             # Forward
  335.             with amp.autocast(enabled=cuda):
  336.                 pred = model(imgs)  # forward
  337.                 loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
  338.                 if RANK != -1:
  339.                     loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
  340.                 if opt.quad:
  341.                     loss *= 4.
  342.  
  343.             # Backward
  344.             scaler.scale(loss).backward()
  345.  
  346.             # Optimize
  347.             if ni - last_opt_step >= accumulate:
  348.                 scaler.step(optimizer)  # optimizer.step
  349.                 scaler.update()
  350.                 optimizer.zero_grad()
  351.                 if ema:
  352.                     ema.update(model)
  353.                 last_opt_step = ni
  354.  
  355.             # Print
  356.             if RANK in [-1, 0]:
  357.                 mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
  358.                 mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
  359.                 s = ('%10s' * 2 + '%10.4g' * 6) % (
  360.                     f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1])
  361.                 pbar.set_description(s)
  362.  
  363.                 # Plot
  364.                 if plots and ni < 3:
  365.                     f = save_dir / f'train_batch{ni}.jpg'  # filename
  366.                     Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
  367.                     if loggers['tb'] and ni == 0:  # TensorBoard
  368.                         with warnings.catch_warnings():
  369.                             warnings.simplefilter('ignore')  # suppress jit trace warning
  370.                             loggers['tb'].add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
  371.                 elif plots and ni == 10 and loggers['wandb']:
  372.                     wandb_logger.log({'Mosaics': [loggers['wandb'].Image(str(x), caption=x.name) for x in
  373.                                                   save_dir.glob('train*.jpg') if x.exists()]})
  374.  
  375.             # end batch ------------------------------------------------------------------------------------------------
  376.  
  377.         # Scheduler
  378.         lr = [x['lr'] for x in optimizer.param_groups]  # for loggers
  379.         scheduler.step()
  380.  
  381.         # DDP process 0 or single-GPU
  382.         if RANK in [-1, 0]:
  383.             # mAP
  384.             ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
  385.             final_epoch = epoch + 1 == epochs
  386.             if not noval or final_epoch:  # Calculate mAP
  387.                 wandb_logger.current_epoch = epoch + 1
  388.                 results, maps, _ = val.run(data_dict,
  389.                                            batch_size=batch_size // WORLD_SIZE * 2,
  390.                                            imgsz=imgsz,
  391.                                            model=ema.ema,
  392.                                            single_cls=single_cls,
  393.                                            dataloader=val_loader,
  394.                                            save_dir=save_dir,
  395.                                            save_json=is_coco and final_epoch,
  396.                                            verbose=nc < 50 and final_epoch,
  397.                                            plots=plots and final_epoch,
  398.                                            wandb_logger=wandb_logger,
  399.                                            compute_loss=compute_loss)
  400.  
  401.             # Write
  402.             with open(results_file, 'a') as f:
  403.                 f.write(s + '%10.4g' * 7 % results + '\n')  # append metrics, val_loss
  404.  
  405.             # Log
  406.             tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss',  # train loss
  407.                     'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
  408.                     'val/box_loss', 'val/obj_loss', 'val/cls_loss',  # val loss
  409.                     'x/lr0', 'x/lr1', 'x/lr2']  # params
  410.             for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
  411.                 if loggers['tb']:
  412.                     loggers['tb'].add_scalar(tag, x, epoch)  # TensorBoard
  413.                 if loggers['wandb']:
  414.                     wandb_logger.log({tag: x})  # W&B
  415.  
  416.             # Update best mAP
  417.             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
  418.             if fi > best_fitness:
  419.                 best_fitness = fi
  420.             wandb_logger.end_epoch(best_result=best_fitness == fi)
  421.  
  422.             # Save model
  423.             if (not nosave) or (final_epoch and not evolve):  # if save
  424.                 ckpt = {'epoch': epoch,
  425.                         'best_fitness': best_fitness,
  426.                         'training_results': results_file.read_text(),
  427.                         'model': deepcopy(de_parallel(model)).half(),
  428.                         'ema': deepcopy(ema.ema).half(),
  429.                         'updates': ema.updates,
  430.                         'optimizer': optimizer.state_dict(),
  431.                         'wandb_id': wandb_logger.wandb_run.id if loggers['wandb'] else None}
  432.  
  433.                 # Save last, best and delete
  434.                 torch.save(ckpt, last)
  435.                 if best_fitness == fi:
  436.                     torch.save(ckpt, best)
  437.                 if loggers['wandb']:
  438.                     if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1:
  439.                         wandb_logger.log_model(last.parent, opt, epoch, fi, best_model=best_fitness == fi)
  440.                 del ckpt
  441.  
  442.         # end epoch ----------------------------------------------------------------------------------------------------
  443.     # end training -----------------------------------------------------------------------------------------------------
  444.     if RANK in [-1, 0]:
  445.         LOGGER.info(f'{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.\n')
  446.         if plots:
  447.             plot_results(save_dir=save_dir)  # save as results.png
  448.             if loggers['wandb']:
  449.                 files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
  450.                 wandb_logger.log({"Results": [loggers['wandb'].Image(str(save_dir / f), caption=f) for f in files
  451.                                               if (save_dir / f).exists()]})
  452.  
  453.         if not evolve:
  454.             if is_coco:  # COCO dataset
  455.                 for m in [last, best] if best.exists() else [last]:  # speed, mAP tests
  456.                     results, _, _ = val.run(data_dict,
  457.                                             batch_size=batch_size // WORLD_SIZE * 2,
  458.                                             imgsz=imgsz,
  459.                                             model=attempt_load(m, device).half(),
  460.                                             single_cls=single_cls,
  461.                                             dataloader=val_loader,
  462.                                             save_dir=save_dir,
  463.                                             save_json=True,
  464.                                             plots=False)
  465.  
  466.             # Strip optimizers
  467.             for f in last, best:
  468.                 if f.exists():
  469.                     strip_optimizer(f)  # strip optimizers
  470.             if loggers['wandb']:  # Log the stripped model
  471.                 loggers['wandb'].log_artifact(str(best if best.exists() else last), type='model',
  472.                                               name='run_' + wandb_logger.wandb_run.id + '_model',
  473.                                               aliases=['latest', 'best', 'stripped'])
  474.         wandb_logger.finish_run()
  475.  
  476.     torch.cuda.empty_cache()
  477.     return results

 

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/很楠不爱3/article/detail/718390
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号