目录
效果展示（完整源码在B站视频简介内）
基于YOLOv5的实时水果识别系统与分类系统演示与介绍（Python+ui界面+训练代码+性能优化+实时检测识别）_哔哩哔哩_bilibili
基于YOLOv5的实时水果识别系统与分类系统演示与介绍（Python+ui界面+训练代码+性能优化+实时检测识别）
（完整源码在B站视频简介内）
环境安装
runs文件夹中，存放训练和评估的结果图
环境安装：
请按照给定的python版本配置环境，否则可能会因依赖不兼容而出错，
在文件目录下cmd进入终端
(1)使用anaconda新建python3.10环境:
conda create -n env_rec python=3.10
(2)激活创建的环境:
conda activate env_rec
(3)使用pip安装所需的依赖，可通过requirements.txt:
pip install -r requirements.txt
在settings中找到project python interpreter 点击Add Interpreter
点击conda，在Use existing environment中选择刚才创建的虚拟环境 ，最后点击确定。如果conda Executable中路径没有，那就把anaconda3的路径添加上
本文介绍了基于Yolov5模型的水果分类识别系统以及使用PyQt库构建的交互式界面。
首先，Yolov5是一种目标检测算法，它可以通过输入图片，自动识别出其中包含的不同目标，并标注出它们的位置和类别。我们利用Yolov5模型对水果图片进行训练，得到一个可靠的水果分类识别系统。该系统准确率高，能够快速地将水果图片分类。
其次，为了更好地与用户进行交互，我们使用PyQt库构建了一个交互式界面。该界面包括了图片上传、执行识别、结果展示等功能模块。用户可以通过上传自己的水果图片，执行识别操作，从而得到自己所拍摄水果的分类结果。界面友好、操作简便，让人可以轻松地使用该系统。
总体而言，这个水果分类识别系统基于Yolov5模型实现，借助PyQt库构建了一个易用性较高的交互式界面。该系统准确率高、速度快、操作简单，是一款非常实用的水果分类识别工具。
Yolov5 是一种先进的目标检测算法，可以应用于水果分类识别任务。结合 PyQT框架，可以创建一个交互式界面，使用户能够方便地上传图片并获取水果分类结果。以下将详细阐述 Yolov5 水果分类识别和 PyQT交互式界面的实现。
Yolov5 是由 Ultralytics公司开发的一种基于深度学习的目标检测算法，它采用了一种称为单阶段目标检测的方法，具有高准确率和实时性的特点。在水果分类识别任务中，Yolov5可以检测图像中的水果，并将其分类为不同的类别，例如苹果、香蕉、橙子等。
为了实现 Yolov5 水果分类识别的交互式界面，可以使用 PyQT 框架进行开发。PyQT 是一个功能强大且易于使用的 PythonGUI 开发工具包，它提供了丰富的界面组件和布局选项，可以轻松创建用户友好的界面。
在界面设计方面，可以使用 PyQT 创建一个包含上传图片按钮和显示分类结果的窗口。当用户点击上传图片按钮时，可以调用 Yolov5模型对上传的图片进行识别，并将分类结果显示在界面上。同时，还可以添加其他功能，如清除界面、保存结果等。
核心源码（完整源码在B站视频简介内）
# -*- coding: UTF-8 -*-
import random
import sys
import threading
import time
import cv2
import numpy
import torch
import torch.backends.cudnn as cudnn
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from models.experimental import attempt_load
from utils.datasets import LoadImages, LoadStreams
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.plots import plot_one_box
from utils.torch_utils import select_device, time_synchronized
model_path = 'runs/train/yolov5s/weights/best.pt'
# 添加一个关于界面
# 窗口主类
class MainWindow(QTabWidget):
# 基本配置不动，然后只动第三个界面
def __init__(self):
# 初始化界面
super().__init__()
self.setWindowTitle('Yolov5水果检测系统')
self.resize(1200, 800)
self.setWindowIcon(QIcon("./UI/xf.jpg"))
# 图片读取进程
self.output_size = 480
self.img2predict = ""
# 空字符串会自己进行选择，首选cuda
self.device = ''
# # 初始化视频读取线程
self.vid_source = '0' # 初始设置为摄像头
# 检测视频的线程
self.threading = None
# 是否跳出当前循环的线程
self.jump_threading: bool = False
self.image_size = 640
self.confidence = 0.25
self.iou_threshold = 0.45
# 指明模型加载的位置的设备
self.model = self.model_load(weights=model_path,
device=self.device)
self.initUI()
self.reset_vid()
@torch.no_grad()
def model_load(self,
weights="", # model.pt path(s)
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
):
"""
模型初始化
"""
device = self.device = select_device(device)
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, device) # load FP32 model
self.stride = int(model.stride.max()) # model stride
self.image_size = check_img_size(self.image_size, s=self.stride) # check img_size
if half:
model.half() # to FP16
# Run inference
if device.type != 'cpu':
print("Run inference")
model(torch.zeros(1, 3, self.image_size, self.image_size).to(device).type_as(
next(model.parameters()))) # run once
print("模型加载完成!")
return model
def reset_vid(self):
"""
界面重置事件
"""
self.webcam_detection_btn.setEnabled(True)
self.mp4_detection_btn.setEnabled(True)
self.left_vid_img.setPixmap(QPixmap("./UI/up.jpeg"))
self.vid_source = '0'
self.disable_btn(self.det_img_button)
self.disable_btn(self.vid_start_stop_btn)
self.jump_threading = False
def initUI(self):
"""
界面初始化
"""
# 图片检测子界面
font_title = QFont('楷体', 16)
font_main = QFont('楷体', 14)
font_general = QFont('楷体', 10)
# 图片识别界面, 两个按钮，上传图片和显示结果
img_detection_widget = QWidget()
img_detection_layout = QVBoxLayout()
img_detection_title = QLabel("图片识别功能")
img_detection_title.setFont(font_title)
mid_img_widget = QWidget()
mid_img_layout = QHBoxLayout()
self.left_img = QLabel()
self.right_img = QLabel()
self.left_img.setPixmap(QPixmap("./UI/up.jpeg"))
self.right_img.setPixmap(QPixmap("./UI/right.jpeg"))
self.left_img.setAlignment(Qt.AlignCenter)
self.right_img.setAlignment(Qt.AlignCenter)
self.left_img.setMinimumSize(480, 480)
self.left_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
mid_img_layout.addWidget(self.left_img)
self.right_img.setMinimumSize(480, 480)
self.right_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
mid_img_layout.addStretch(0)
mid_img_layout.addWidget(self.right_img)
mid_img_widget.setLayout(mid_img_layout)
self.up_img_button = QPushButton("上传图片")
self.det_img_button = QPushButton("开始检测")
self.up_img_button.clicked.connect(self.upload_img)
self.det_img_button.clicked.connect(self.detect_img)
self.up_img_button.setFont(font_main)
self.det_img_button.setFont(font_main)
self.up_img_button.setStyleSheet("QPushButton{color:white}"
"QPushButton:hover{background-color: rgb(2,110,180);}"
"QPushButton{background-color:rgb(48,124,208)}"
"QPushButton{border:2px}"
"QPushButton{border-radius:5px}"
"QPushButton{padding:5px 5px}"
"QPushButton{margin:5px 5px}")
self.det_img_button.setStyleSheet("QPushButton{color:white}"
"QPushButton:hover{background-color: rgb(2,110,180);}"
"QPushButton{background-color:rgb(48,124,208)}"
"QPushButton{border:2px}"
"QPushButton{border-radius:5px}"
"QPushButton{padding:5px 5px}"
"QPushButton{margin:5px 5px}")
img_detection_layout.addWidget(img_detection_title, alignment=Qt.AlignCenter)
img_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
img_detection_layout.addWidget(self.up_img_button)
img_detection_layout.addWidget(self.det_img_button)
img_detection_widget.setLayout(img_detection_layout)
# 视频识别界面
# 视频识别界面的逻辑比较简单，基本就从上到下的逻辑
vid_detection_widget = QWidget()
vid_detection_layout = QVBoxLayout()
vid_title = QLabel("视频检测功能")
vid_title.setFont(font_title)
self.left_vid_img = QLabel()
self.right_vid_img = QLabel()
self.left_vid_img.setPixmap(QPixmap("./UI/up.jpeg"))
self.right_vid_img.setPixmap(QPixmap("./UI/right.jpeg"))
self.left_vid_img.setAlignment(Qt.AlignCenter)
self.left_vid_img.setMinimumSize(480, 480)
self.left_vid_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
self.right_vid_img.setAlignment(Qt.AlignCenter)
self.right_vid_img.setMinimumSize(480, 480)
self.right_vid_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
mid_img_widget = QWidget()
mid_img_layout = QHBoxLayout()
mid_img_layout.addWidget(self.left_vid_img)
mid_img_layout.addStretch(0)
mid_img_layout.addWidget(self.right_vid_img)
mid_img_widget.setLayout(mid_img_layout)
self.webcam_detection_btn = QPushButton("摄像头实时监测")
self.mp4_detection_btn = QPushButton("视频文件检测")
self.vid_start_stop_btn = QPushButton("启动/停止检测")
self.webcam_detection_btn.setFont(font_main)
self.mp4_detection_btn.setFont(font_main)
self.vid_start_stop_btn.setFont(font_main)
self.webcam_detection_btn.setStyleSheet("QPushButton{color:white}"
"QPushButton:hover{background-color: rgb(2,110,180);}"
"QPushButton{background-color:rgb(48,124,208)}"
"QPushButton{border:2px}"
"QPushButton{border-radius:5px}"
"QPushButton{padding:5px 5px}"
"QPushButton{margin:5px 5px}")
self.mp4_detection_btn.setStyleSheet("QPushButton{color:white}"
"QPushButton:hover{background-color: rgb(2,110,180);}"
"QPushButton{background-color:rgb(48,124,208)}"
"QPushButton{border:1px}"
"QPushButton{border-radius:5px}"
"QPushButton{padding:5px 5px}"
"QPushButton{margin:5px 5px}")
self.vid_start_stop_btn.setStyleSheet("QPushButton{color:white}"
"QPushButton:hover{background-color: rgb(2,110,180);}"
"QPushButton{background-color:rgb(48,124,208)}"
"QPushButton{border:2px}"
"QPushButton{border-radius:5px}"
"QPushButton{padding:5px 5px}"
"QPushButton{margin:5px 5px}")
self.webcam_detection_btn.clicked.connect(self.open_cam)
self.mp4_detection_btn.clicked.connect(self.open_mp4)
self.vid_start_stop_btn.clicked.connect(self.start_or_stop)
# 添加fps显示
fps_container = QWidget()
fps_container.setStyleSheet("QWidget{background-color: #f6f8fa;}")
fps_container_layout = QHBoxLayout()
fps_container.setLayout(fps_container_layout)
# 左容器
fps_left_container = QWidget()
fps_left_container.setStyleSheet("QWidget{background-color: #f6f8fa;}")
fps_left_container_layout = QHBoxLayout()
fps_left_container.setLayout(fps_left_container_layout)
# 右容器
fps_right_container = QWidget()
fps_right_container.setStyleSheet("QWidget{background-color: #f6f8fa;}")
fps_right_container_layout = QHBoxLayout()
fps_right_container.setLayout(fps_right_container_layout)
# 将左容器和右容器添加到fps_container_layout中
fps_container_layout.addWidget(fps_left_container)
fps_container_layout.addStretch(0)
fps_container_layout.addWidget(fps_right_container)
# 左容器中添加fps显示
raw_fps_label = QLabel("原始帧率:")
raw_fps_label.setFont(font_general)
raw_fps_label.setAlignment(Qt.AlignLeft)
raw_fps_label.setStyleSheet("QLabel{margin-left:80px}")
self.raw_fps_value = QLabel("0")
self.raw_fps_value.setFont(font_general)
self.raw_fps_value.setAlignment(Qt.AlignLeft)
fps_left_container_layout.addWidget(raw_fps_label)
fps_left_container_layout.addWidget(self.raw_fps_value)
# 右容器中添加fps显示
detect_fps_label = QLabel("检测帧率:")
detect_fps_label.setFont(font_general)
detect_fps_label.setAlignment(Qt.AlignRight)
self.detect_fps_value = QLabel("0")
self.detect_fps_value.setFont(font_general)
self.detect_fps_value.setAlignment(Qt.AlignRight)
self.detect_fps_value.setStyleSheet("QLabel{margin-right:96px}")
fps_right_container_layout.addWidget(detect_fps_label)
fps_right_container_layout.addWidget(self.detect_fps_value)
# 添加组件到布局上
vid_detection_layout.addWidget(vid_title, alignment=Qt.AlignCenter)
vid_detection_layout.addWidget(fps_container)
vid_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
vid_detection_layout.addWidget(self.webcam_detection_btn)
vid_detection_layout.addWidget(self.mp4_detection_btn)
vid_detection_layout.addWidget(self.vid_start_stop_btn)
vid_detection_widget.setLayout(vid_detection_layout)
# 关于界面
about_widget = QWidget()
about_layout = QVBoxLayout()
about_title = QLabel('CSDN：https://blog.csdn.net/m0_68036862?type=blog\n哔哩哔哩：https://space.bilibili.com/1587012148?spm_id_from=333.1007.0.0\nCSDN搜索人工智能_SYBH') # 修改欢迎词语
about_title.setFont(QFont('楷体', 18))
about_title.setAlignment(Qt.AlignCenter)
about_img = QLabel()
about_img.setPixmap(QPixmap('./UI/qq.png'))
about_img.setAlignment(Qt.AlignCenter)
label_super = QLabel() # 更换作者信息
label_super.setText("人工智能_SYBH")
label_super.setFont(QFont('楷体', 16))
label_super.setOpenExternalLinks(True)
label_super.setOpenExternalLinks(True)
label_super.setAlignment(Qt.AlignRight)
about_layout.addWidget(about_title)
about_layout.addStretch()
about_layout.addWidget(about_img)
about_layout.addStretch()
about_layout.addWidget(label_super)
about_widget.setLayout(about_layout)
self.addTab(img_detection_widget, '图片检测')
self.addTab(vid_detection_widget, '视频检测')
self.addTab(about_widget, '问题咨询')
self.setTabIcon(0, QIcon('./UI/lufei.png'))
self.setTabIcon(1, QIcon('./UI/lufei.png'))
def disable_btn(self, pushButton: QPushButton):
pushButton.setDisabled(True)
pushButton.setStyleSheet("QPushButton{background-color: rgb(2,110,180);}")
def enable_btn(self, pushButton: QPushButton):
pushButton.setEnabled(True)
pushButton.setStyleSheet(
"QPushButton{background-color: rgb(48,124,208);}"
"QPushButton{color:white}"
)
def detect(self, source: str, left_img: QLabel, right_img: QLabel):
"""
@param source: file/dir/URL/glob, 0 for webcam
@param left_img: 将左侧QLabel对象传入，用于显示图片
@param right_img: 将右侧QLabel对象传入，用于显示图片
"""
model = self.model
img_size = [self.image_size, self.image_size] # inference size (pixels)
conf_threshold = self.confidence # confidence threshold
iou_threshold = self.iou_threshold # NMS IOU threshold
device = self.device # cuda device, i.e. 0 or 0,1,2,3 or cpu
classes = None # filter by class: --class 0, or --class 0 2 3
agnostic_nms = False # class-agnostic NMS
augment = False # augmented inference
half = device.type != 'cpu' # half precision only supported on CUDA
if source == "":
self.disable_btn(self.det_img_button)
QMessageBox.warning(self, "请上传", "请先上传视频或图片再进行检测")
else:
source = str(source)
webcam = source.isnumeric()
# Set Dataloader
if webcam:
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=img_size, stride=self.stride)
else:
dataset = LoadImages(source, img_size=img_size, stride=self.stride)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
# 用来记录处理的图片数量
count = 0
# 计算帧率开始时间
fps_start_time = time.time()
for path, img, im0s, vid_cap in dataset:
# 直接跳出for，结束线程
if self.jump_threading:
# 清除状态
self.jump_threading = False
break
count += 1
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=augment)[0]
# Apply NMS
pred = non_max_suppression(pred, conf_threshold, iou_threshold, classes=classes, agnostic=agnostic_nms)
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
s, im0 = 'detect : ', im0s[i].copy()
else:
s, im0 = 'detect : ', im0s.copy()
# s += '%gx%g ' % img.shape[2:] # print string
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
if webcam or vid_cap is not None:
if webcam: # batch_size >= 1
img = im0s[i]
else:
img = im0s
img = self.resize_img(img)
img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1],
QImage.Format_RGB888)
left_img.setPixmap(QPixmap.fromImage(img))
# 计算一次帧率
if count % 10 == 0:
fps = int(10 / (time.time() - fps_start_time))
self.detect_fps_value.setText(str(fps))
fps_start_time = time.time()
# 应该调整一下图片的大小
# 时间显示
timenumber = time.strftime('%Y/%m/%d/-%H:%M:%S', time.localtime(time.time()))
im0 = cv2.putText(im0, timenumber, (50, 50), cv2.FONT_HERSHEY_SIMPLEX,
1, (0, 255, 0), 2, cv2.LINE_AA)
im0 = cv2.putText(im0, s, (50, 80), cv2.FONT_HERSHEY_SIMPLEX,
1, (255, 0, 0), 2, cv2.LINE_AA)
img = self.resize_img(im0)
img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1],
QImage.Format_RGB888)
right_img.setPixmap(QPixmap.fromImage(img))
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
# 使用完摄像头释放资源
if webcam:
for cap in dataset.caps:
cap.release()
else:
dataset.cap and dataset.cap.release()
def resize_img(self, img):
"""
调整图片大小，方便用来显示
@param img: 需要调整的图片
"""
resize_scale = min(self.output_size / img.shape[0], self.output_size / img.shape[1])
img = cv2.resize(img, (0, 0), fx=resize_scale, fy=resize_scale)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
def upload_img(self):
"""
上传图片
"""
# 选择录像文件进行读取
fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.jpg *.png *.tif *.jpeg')
if fileName:
self.img2predict = fileName
# 将上传照片和进行检测做成互斥的
self.enable_btn(self.det_img_button)
self.disable_btn(self.up_img_button)
# 进行左侧原图展示
img = cv2.imread(fileName)
# 应该调整一下图片的大小
img = self.resize_img(img)
img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1], QImage.Format_RGB888)
self.left_img.setPixmap(QPixmap.fromImage(img))
# 上传图片之后右侧的图片重置
self.right_img.setPixmap(QPixmap("./UI/right.jpeg"))
def detect_img(self):
"""
检测图片
"""
# 重置跳出线程状态，防止其他位置使用的影响
self.jump_threading = False
self.detect(self.img2predict, self.left_img, self.right_img)
# 将上传照片和进行检测做成互斥的
self.enable_btn(self.up_img_button)
self.disable_btn(self.det_img_button)
def open_mp4(self):
"""
开启视频文件检测事件
"""
print("开启视频文件检测")
fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.mp4 *.avi')
if fileName:
self.disable_btn(self.webcam_detection_btn)
self.disable_btn(self.mp4_detection_btn)
self.enable_btn(self.vid_start_stop_btn)
# 生成读取视频对象
cap = cv2.VideoCapture(fileName)
# 获取视频的帧率
fps = cap.get(cv2.CAP_PROP_FPS)
# 显示原始视频帧率
self.raw_fps_value.setText(str(fps))
if cap.isOpened():
# 读取一帧用来提前左侧展示
ret, raw_img = cap.read()
cap.release()
else:
QMessageBox.warning(self, "需要重新上传", "请重新选择视频文件")
self.disable_btn(self.vid_start_stop_btn)
self.enable_btn(self.webcam_detection_btn)
self.enable_btn(self.mp4_detection_btn)
return
# 应该调整一下图片的大小
img = self.resize_img(numpy.array(raw_img))
img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1], QImage.Format_RGB888)
self.left_vid_img.setPixmap(QPixmap.fromImage(img))
# 上传图片之后右侧的图片重置
self.right_vid_img.setPixmap(QPixmap("./UI/right.jpeg"))
self.vid_source = fileName
self.jump_threading = False
def open_cam(self):
"""
打开摄像头事件
"""
print("打开摄像头")
self.disable_btn(self.webcam_detection_btn)
self.disable_btn(self.mp4_detection_btn)
self.enable_btn(self.vid_start_stop_btn)
self.vid_source = "0"
self.jump_threading = False
# 生成读取视频对象
cap = cv2.VideoCapture(0)
# 获取视频的帧率
fps = cap.get(cv2.CAP_PROP_FPS)
# 显示原始视频帧率
self.raw_fps_value.setText(str(fps))
if cap.isOpened():
# 读取一帧用来提前左侧展示
ret, raw_img = cap.read()
cap.release()
else:
QMessageBox.warning(self, "需要重新上传", "请重新选择视频文件")
self.disable_btn(self.vid_start_stop_btn)
self.enable_btn(self.webcam_detection_btn)
self.enable_btn(self.mp4_detection_btn)
return
# 应该调整一下图片的大小
img = self.resize_img(numpy.array(raw_img))
img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1], QImage.Format_RGB888)
self.left_vid_img.setPixmap(QPixmap.fromImage(img))
# 上传图片之后右侧的图片重置
self.right_vid_img.setPixmap(QPixmap("./UI/right.jpeg"))
def start_or_stop(self):
"""
启动或者停止事件
"""
print("启动或者停止")
if self.threading is None:
# 创造并启动一个检测视频线程
self.jump_threading = False
self.threading = threading.Thread(target=self.detect_vid)
self.threading.start()
self.disable_btn(self.webcam_detection_btn)
self.disable_btn(self.mp4_detection_btn)
else:
# 停止当前线程
# 线程属性置空，恢复状态
self.threading = None
self.jump_threading = True
self.enable_btn(self.webcam_detection_btn)
self.enable_btn(self.mp4_detection_btn)
def detect_vid(self):
"""
视频检测
视频和摄像头的主函数是一样的，不过是传入的source不同罢了
"""
print("视频开始检测")
self.detect(self.vid_source, self.left_vid_img, self.right_vid_img)
print("视频检测结束")
# 执行完进程，刷新一下和进程有关的状态，只有self.threading是None，
# 才能说明是正常结束的线程，需要被刷新状态
if self.threading is not None:
self.start_or_stop()
def closeEvent(self, event):
"""
界面关闭事件
"""
reply = QMessageBox.question(
self,
'quit',
"Are you sure?",
QMessageBox.Yes | QMessageBox.No,
QMessageBox.No
)
if reply == QMessageBox.Yes:
self.jump_threading = True
self.close()
event.accept()
else:
event.ignore()
if __name__ == "__main__":
app = QApplication(sys.argv)
mainWindow = MainWindow()
mainWindow.show()
sys.exit(app.exec_())
train.py（完整源码在B站视频简介内）
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Train a YOLOv5 model on a custom dataset.
Models and datasets download automatically from the latest YOLOv5 release.
Usage - Single-GPU training:
$ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended)
$ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch
Usage - Multi-GPU DDP training:
$ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,3
Models: https://github.com/ultralytics/yolov5/tree/master/models
Datasets: https://github.com/ultralytics/yolov5/tree/master/data
Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
"""
import argparse
import math
import os
import random
import sys
import time
from copy import deepcopy
from datetime import datetime
from pathlib import Path
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
import val as validate # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.downloads import attempt_download, is_url
from utils.general import (LOGGER, check_amp, check_dataset, check_file, check_git_status, check_img_size,
check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path,
init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods,
one_cycle, print_args, print_mutation, strip_optimizer, yaml_save)
from utils.loggers import Loggers
from utils.loggers.comet.comet_utils import check_comet_resume
from utils.loggers.wandb.wandb_utils import check_wandb_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve
from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
smart_resume, torch_distributed_zero_first)
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv('RANK', -1))
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
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
callbacks.run('on_pretrain_routine_start')
# Directories
w = save_dir / 'weights' # weights dir
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
last, best = w / 'last.pt', w / 'best.pt'
# Hyperparameters
if isinstance(hyp, str):
with open(hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
opt.hyp = hyp.copy() # for saving hyps to checkpoints
# Save run settings
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp)
yaml_save(save_dir / 'opt.yaml', vars(opt))
# Loggers
data_dict = None
if RANK in {-1, 0}:
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
# Register actions
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))
# Process custom dataset artifact link
data_dict = loggers.remote_dataset
if resume: # If resuming runs from remote artifact
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
# Config
plots = not evolve and not opt.noplots # create plots
cuda = device.type != 'cpu'
init_seeds(opt.seed + 1 + RANK, deterministic=True)
with torch_distributed_zero_first(LOCAL_RANK):
data_dict = data_dict or check_dataset(data) # check if None
train_path, val_path = data_dict['train'], data_dict['val']
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
# Model
check_suffix(weights, '.pt') # check weights
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location='cpu') # load checkpoint to CPU to avoid CUDA memory leak
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect
model.load_state_dict(csd, strict=False) # load
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
else:
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
amp = check_amp(model) # check AMP
# Freeze
freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
# v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)
if any(x in k for x in freeze):
LOGGER.info(f'freezing {k}')
v.requires_grad = False
# Image size
gs = max(int(model.stride.max()), 32) # grid size (max stride)
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
# Batch size
if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
batch_size = check_train_batch_size(model, imgsz, amp)
loggers.on_params_update({"batch_size": batch_size})
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
# Scheduler
if opt.cos_lr:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
else:
lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA
ema = ModelEMA(model) if RANK in {-1, 0} else None
# Resume
best_fitness, start_epoch = 0.0, 0
if pretrained:
if resume:
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
del ckpt, csd
# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
# Trainloader
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect,
rank=LOCAL_RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
shuffle=True)
labels = np.concatenate(dataset.labels, 0)
mlc = int(labels[:, 0].max()) # max label class
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# Process 0
if RANK in {-1, 0}:
val_loader = create_dataloader(val_path,
imgsz,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=None if noval else opt.cache,
rect=True,
rank=-1,
workers=workers * 2,
pad=0.5,
prefix=colorstr('val: '))[0]
if not resume:
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor
model.half().float() # pre-reduce anchor precision
callbacks.run('on_pretrain_routine_end', labels, names)
# DDP mode
if cuda and RANK != -1:
model = smart_DDP(model)
# Model attributes
nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)
hyp['box'] *= 3 / nl # scale to layers
hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Start training
t0 = time.time()
nb = len(train_loader) # number of batches
nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
last_opt_step = -1
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = torch.cuda.amp.GradScaler(enabled=amp)
stopper, stop = EarlyStopping(patience=opt.patience), False
compute_loss = ComputeLoss(model) # init loss class
callbacks.run('on_train_start')
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
f"Logging results to {colorstr('bold', save_dir)}\n"
f'Starting training for {epochs} epochs...')
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
callbacks.run('on_train_epoch_start')
model.train()
# Update image weights (optional, single-GPU only)
if opt.image_weights:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
# Update mosaic border (optional)
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(3, device=device) # mean losses
if RANK != -1:
train_loader.sampler.set_epoch(epoch)
pbar = enumerate(train_loader)
LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
if RANK in {-1, 0}:
pbar = tqdm(pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
callbacks.run('on_train_batch_start')
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
with torch.cuda.amp.autocast(amp):
pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
if ni - last_opt_step >= accumulate:
scaler.unscale_(optimizer) # unscale gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
last_opt_step = ni
# Log
if RANK in {-1, 0}:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
if callbacks.stop_training:
return
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
if RANK in {-1, 0}:
# mAP
callbacks.run('on_train_epoch_end', epoch=epoch)
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
if not noval or final_epoch: # Calculate mAP
results, maps, _ = validate.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
half=amp,
model=ema.ema,
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
plots=False,
callbacks=callbacks,
compute_loss=compute_loss)
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
stop = stopper(epoch=epoch, fitness=fi) # early stop check
if fi > best_fitness:
best_fitness = fi
log_vals = list(mloss) + list(results) + lr
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,
'opt': vars(opt),
'date': datetime.now().isoformat()}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if opt.save_period > 0 and epoch % opt.save_period == 0:
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
# EarlyStopping
if RANK != -1: # if DDP training
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # must break all DDP ranks
# end epoch ----------------------------------------------------------------------------------------------------
# end training -----------------------------------------------------------------------------------------------------
if RANK in {-1, 0}:
LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
for f in last, best:
if f.exists():
strip_optimizer(f) # strip optimizers
if f is best:
LOGGER.info(f'\nValidating {f}...')
results, _, _ = validate.run(
data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
model=attempt_load(f, device).half(),
iou_thres=0.65 if is_coco else 0.60, # best pycocotools at iou 0.65
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=plots,
callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots
if is_coco:
callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
callbacks.run('on_train_end', last, best, epoch, results)
torch.cuda.empty_cache()
return results
def parse_opt(known=False):
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default=ROOT / 'weights/yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='models/ShuffleNetV2.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/VOC.yaml', help='dataset.yaml path')
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs, -1 for autobatch')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--noval', action='store_true', help='only validate final epoch')
parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
parser.add_argument('--noplots', action='store_true', help='save no plot files')
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--workers', type=int, default=2, help='max dataloader workers (per RANK in DDP mode)')
parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
parser.add_argument('--name', default='exp', help='save to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--quad', action='store_true', help='quad dataloader')
parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
parser.add_argument('--seed', type=int, default=0, help='Global training seed')
parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
# Logger arguments
parser.add_argument('--entity', default=None, help='Entity')
parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
return parser.parse_known_args()[0] if known else parser.parse_args()
def main(opt, callbacks=Callbacks()):
# Checks
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
check_requirements()
# Resume (from specified or most recent last.pt)
if opt.resume and not check_wandb_resume(opt) and not check_comet_resume(opt) and not opt.evolve:
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml
opt_data = opt.data # original dataset
if opt_yaml.is_file():
with open(opt_yaml, errors='ignore') as f:
d = yaml.safe_load(f)
else:
d = torch.load(last, map_location='cpu')['opt']
opt = argparse.Namespace(**d) # replace
opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate
if is_url(opt_data):
opt.data = check_file(opt_data) # avoid HUB resume auth timeout
else:
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
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
if opt.evolve:
if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve
opt.project = str(ROOT / 'runs/evolve')
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
if opt.name == 'cfg':
opt.name = Path(opt.cfg).stem # use model.yaml as name
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
# DDP mode
device = select_device(opt.device, batch_size=opt.batch_size)
if LOCAL_RANK != -1:
msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
assert not opt.image_weights, f'--image-weights {msg}'
assert not opt.evolve, f'--evolve {msg}'
assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
torch.cuda.set_device(LOCAL_RANK)
device = torch.device('cuda', LOCAL_RANK)
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
# Train
if not opt.evolve:
train(opt.hyp, opt, device, callbacks)
# Evolve hyperparameters (optional)
else:
# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
meta = {
'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'box': (1, 0.02, 0.2), # box loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0), # image mixup (probability)
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
with open(opt.hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict
if 'anchors' not in hyp: # anchors commented in hyp.yaml
hyp['anchors'] = 3
if opt.noautoanchor:
del hyp['anchors'], meta['anchors']
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
if opt.bucket:
os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}') # download evolve.csv if exists
for _ in range(opt.evolve): # generations to evolve
if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
# Select parent(s)
parent = 'single' # parent selection method: 'single' or 'weighted'
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
n = min(5, len(x)) # number of previous results to consider
x = x[np.argsort(-fitness(x))][:n] # top n mutations
w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
if parent == 'single' or len(x) == 1:
# x = x[random.randint(0, n - 1)] # random selection
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
elif parent == 'weighted':
x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
# Mutate
mp, s = 0.8, 0.2 # mutation probability, sigma
npr = np.random
npr.seed(int(time.time()))
g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
ng = len(meta)
v = np.ones(ng)
while all(v == 1): # mutate until a change occurs (prevent duplicates)
v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# Constrain to limits
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # lower limit
hyp[k] = min(hyp[k], v[2]) # upper limit
hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation
results = train(hyp.copy(), opt, device, callbacks)
callbacks = Callbacks()
# Write mutation results
print_mutation(results, hyp.copy(), save_dir, opt.bucket)
# Plot results
plot_evolve(evolve_csv)
LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
f"Results saved to {colorstr('bold', save_dir)}\n"
f'Usage example: $ python train.py --hyp {evolve_yaml}')
def run(**kwargs):
# Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
opt = parse_opt(True)
for k, v in kwargs.items():
setattr(opt, k, v)
main(opt)
return opt
if __name__ == "__main__":
opt = parse_opt()
main(opt)
扫一扫
3637
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
