Skip to content

Instantly share code, notes, and snippets.

@dyd1234
Created January 23, 2025 13:44
Show Gist options
  • Save dyd1234/f4faad5a5fbf35aaccb1d3bb1a8b630f to your computer and use it in GitHub Desktop.
Save dyd1234/f4faad5a5fbf35aaccb1d3bb1a8b630f to your computer and use it in GitHub Desktop.
Train for YOLO v5s with convering the detection head to linear layer
# Ultralytics YOLOv5 🚀, AGPL-3.0 license
# Hyperparameters for low-augmentation COCO training from scratch
# python train.py --batch 64 --cfg yolov5n6.yaml --weights '' --data coco.yaml --img 640 --epochs 300 --linear
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
lr0: 0.001 # initial learning rate (SGD=1E-2, Adam=1E-3)
lrf: 0.01 # final OneCycleLR learning rate (lr0 * lrf)
# lr0: 0.0002 # initial learning rate (SGD=1E-2, Adam=1E-3)
# lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf)
momentum: 0.937 # SGD momentum/Adam beta1
weight_decay: 0.0005 # optimizer weight decay 5e-4
warmup_epochs: 5.0 # warmup epochs (fractions ok)
warmup_momentum: 0.8 # warmup initial momentum
warmup_bias_lr: 0.1 # warmup initial bias lr
box: 0.05 # box loss gain
cls: 0.5 # cls loss gain
cls_pw: 1.0 # cls BCELoss positive_weight
obj: 1.0 # obj loss gain (scale with pixels)
obj_pw: 1.0 # obj BCELoss positive_weight
iou_t: 0.20 # IoU training threshold
anchor_t: 4.0 # anchor-multiple threshold
# anchors: 3 # anchors per output layer (0 to ignore)
fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
hsv_v: 0.4 # image HSV-Value augmentation (fraction)
degrees: 0.0 # image rotation (+/- deg)
translate: 0.0 # image translation (+/- fraction)
scale: 0.0 # image scale (+/- gain)
shear: 0.0 # image shear (+/- deg)
perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
flipud: 0.0 # image flip up-down (probability)
fliplr: 0.0 # image flip left-right (probability)
mosaic: 0.0 # image mosaic (probability)
mixup: 0.0 # image mixup (probability)
copy_paste: 0.0 # segment copy-paste (probability)
# Ultralytics YOLOv5 🚀, AGPL-3.0 license
# Parameters
nc: 60 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
# - [10, 13, 16, 30, 33, 23] # P3/8
# - [30, 61, 62, 45, 59, 119] # P4/16
- [116, 90, 156, 198, 373, 326] # P5/32
- [116, 90, 156, 198, 373, 326] # P5/32
- [116, 90, 156, 198, 373, 326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[ #
[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# 看看能不能直接从这里开始进行设置
# YOLOv5 v6.0 head
head: [
[-1, 1, Conv, [512, 1, 1]], # 10
# [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 11
# [[-1, -1], 1, Concat, [1]], # cat backbone P4 # 12
# [-1, 3, C3, [512, False]], # 13
# [-1, 1, Conv, [256, 1, 1]], # 14
# [-1, 1, nn.Upsample, [None, 2, "nearest"]],
# [[-1, 4], 1, Concat, [1]], # 16 cat backbone P3
# [-1, 3, C3, [256, False]], # 17 (P3/8-small)
# [-1, 1, Conv, [256, 3, 2]],
# [[-1, 14], 1, Concat, [1]], # cat head P4
# [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
# [-1, 1, Conv, [512, 3, 2]],
[[10, 10], 1, Concat, [1]], # cat head P5 # new 11 22
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
# 直接只从一层取数据
# [[23, 23, 23], 1, Detect, [nc, anchors]], # # Detect(P3, P4, P5)
[[12, 12, 12], 1, Detect, [nc, anchors]], # # Detect(P3, P4, P5)
]
# Ultralytics YOLOv5 🚀, AGPL-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://docs.ultralytics.com/yolov5/tutorials/train_custom_data
"""
import argparse
import math
import os
import random
import subprocess
import sys
import time
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path
import wandb
# try:
# import comet_ml # must be imported before torch (if installed)
# except ImportError:
comet_ml = None
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
from scn_utlis import apply_rotation_augmentation, save_comparison_images, scn_config_praser, transform_angle
from scn import SCN, ScnLayer
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, Detect # add one more Detect class defination
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.segment.dataloaders import create_dataloader
from utils.segment.augmentations import apply_test_rotation
from utils.downloads import attempt_download, is_url
from utils.general import (
LOGGER,
TQDM_BAR_FORMAT,
check_amp,
check_dataset,
check_file,
check_git_info,
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, Loggers
from utils.loggers.comet.comet_utils import check_comet_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,
)
from scn_utlis import print_trainable_modules
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))
GIT_INFO = check_git_info()
cos = nn.CosineSimilarity(dim=0, eps=1e-6)
class LinearLayer(nn.Module):
def __init__(self, in_features, out_features, num_classes, i, f, anchors=(), inplace=True):
super(LinearLayer, self).__init__()
self.flatten = nn.Flatten() # 添加 Flatten 层
self.linear = nn.Linear(in_features, out_features)
# 复制 Detect 类中的相关属性
self.nc = num_classes # number of classes
self.no = num_classes + 5 # number of outputs (可以根据需要调整)
self.nl = 1 # 线性层只有一个输出层
self.na = len(anchors) // 2 if anchors else 0 # number of anchors
self.anchors = torch.tensor(anchors).float().view(self.nl, -1, 2) if anchors else None # shape(nl, na, 2)
self.inplace = inplace # use inplace ops (e.g. slice assignment)
self.i = i
self.f = f
def forward(self, x):
# 检查输入是否为列表
if isinstance(x, list):
if len(x) == 0:
raise ValueError("Input list is empty.")
x = x[0] # 选择列表中的第一个张量
# 检查输入是否为张量
if not isinstance(x, torch.Tensor):
raise TypeError(f"Expected input to be a torch.Tensor, but got {type(x)}")
# # 检查张量是否在 GPU 上
# if x.is_cuda:
# print("Input tensor is on GPU.")
# else:
# print("Input tensor is on CPU.")
# # 打印输入形状
# print(f"Input shape: {x.shape}")
# # 检查输入是否为张量
# if not isinstance(x, torch.Tensor):
# raise TypeError(f"Expected input to be a torch.Tensor, but got {type(x)}")
# # 打印输入形状
# print(f"Input shape: {x.shape}")
x = self.flatten(x) # 在前向传播中先展平输入
return self.linear(x)
def get_output_shape(self, input_shape):
"""根据输入形状返回输出形状"""
# 这里可以根据需要实现输出形状的计算
return (input_shape[0], self.no) # 假设输出形状为 (batch_size, no)
def replace_detect_with_linear(model, num_classes, anchors=()):
for i, layer in enumerate(model.model):
if isinstance(layer, Detect):
# 计算输入特征数
# in_features = 512 * 20 * 20 # 假设特征图大小为 20x20
in_features = 512 * 10 * 10 # deep one with 320*320 input
# in_features = 512 * 20 * 20 # mid one with 320*320 input
# in_features = 512 * 40 * 40 # shallow one with 320*320 input
# some other features
f = model.model[i].f
# print
# 创建 LinearLayer
linear_layer = LinearLayer(in_features, num_classes, num_classes, i, f, anchors=anchors) # 使用 num_classes 和 anchors
# linear_layer.to_device()
linear_layer.to(next(model.parameters()).device) # 将 linear_layer 移动到模型的设备
# 替换 Detect 层
model.model[i] = linear_layer
# 打印替换信息
print(f"Replaced Detect layer with LinearLayer at index {i}.")
print(f"Input features: {in_features}, Output features: {num_classes}")
break # 替换后退出循环
def train(hyp, opt, device, callbacks):
"""
Train a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets,
model architecture, loss computation, and optimizer steps.
Args:
hyp (str | dict): Path to the hyperparameters YAML file or a dictionary of hyperparameters.
opt (argparse.Namespace): Parsed command-line arguments containing training options.
device (torch.device): Device on which training occurs, e.g., 'cuda' or 'cpu'.
callbacks (Callbacks): Callback functions for various training events.
Returns:
None
Models and datasets download automatically from the latest YOLOv5 release.
Example:
Single-GPU training:
```bash
$ 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
```
Multi-GPU DDP training:
```bash
$ 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
```
For more usage details, refer to:
- Models: https://github.com/ultralytics/yolov5/tree/master/models
- Datasets: https://github.com/ultralytics/yolov5/tree/master/data
- Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
"""
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}:
include_loggers = list(LOGGERS)
if getattr(opt, "ndjson_console", False):
include_loggers.append("ndjson_console")
if getattr(opt, "ndjson_file", False):
include_loggers.append("ndjson_file")
loggers = Loggers(
save_dir=save_dir,
weights=weights,
opt=opt,
hyp=hyp,
logger=LOGGER,
include=tuple(include_loggers), #
)
# 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
scn_config = scn_config_praser(opt.scn_config) if opt.scn else None
# 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
# print_trainable_modules(model)
amp = check_amp(model) # check AMP
# print("*"*100)
# print(model)
# print("*"*100)
if opt.scn:
# Yolov5's model located at DetectionMode.model
# model = SCN(num_alpha=scn_config["num_alpha"], dimensions=scn_config["dimensions"], base_model=model, config=scn_config)
if opt.dimensions != -1:
setattr(model, "model", SCN(num_alpha=scn_config["num_alpha"], dimensions=opt.dimensions, base_model=model.model, config=scn_config))
else:
setattr(model, "model", SCN(num_alpha=scn_config["num_alpha"], dimensions=scn_config["dimensions"], base_model=model.model, config=scn_config))
model.to(device)
opt.test_angle = "random"
print(f"The following layers are ScnLayer:")
for name, module in model.named_modules():
if isinstance(module, ScnLayer):
print(f"{name}: {module}")
# replace_detect_with_linear(model, num_classes=nc) # 假设有80个类
# # 打印替换后的模型结构
# print("Updated model structure after replacing Detect with LinearLayer:")
# print(model) # 打印模型结构
# # 如果需要,可以打印每一层的详细信息
# for name, module in model.named_modules():
# print(f"{name}: {module}") # 打印每一层的名称和类型
# input() # 在这里进行打印等待
# 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:
def lf(x):
"""Linear learning rate scheduler function with decay calculated by epoch proportion."""
return (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://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training 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,
seed=opt.seed,
)
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
# align to the training dataset
# if RANK in {-1, 0}:
# val_loader = create_dataloader(
# val_path,
# imgsz,
# batch_size // WORLD_SIZE, #
# gs,
# single_cls,
# hyp=hyp,
# cache=None if noval else opt.cache,
# rect=opt.rect,
# rank=-1,
# workers=workers * 2,
# pad=0.5,
# prefix=colorstr("val: "),
# )[0]
if RANK in {-1, 0}:
val_loader = create_dataloader(
val_path,
imgsz,
batch_size // WORLD_SIZE, # 对齐 batch_size
gs,
single_cls,
hyp=hyp,
cache=None if noval else opt.cache, # 对齐 cache
rect=opt.rect, # 对齐 rect
rank=-1, # 对齐 rank
workers=workers, # 对齐 workers
image_weights=opt.image_weights, # 对齐 image_weights
quad=opt.quad, # 对齐 quad
prefix=colorstr("val: "), # 对齐 prefix
shuffle=False, # 通常验证集不需要打乱
seed=opt.seed, # 对齐 seed
)[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, [email protected], [email protected], 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
# compute_loss = nn.BCEWithLogitsLoss() # For classification model results to make
compute_loss = nn.CrossEntropyLoss() # For multi-class classification
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...'
)
replace_detect_with_linear(model, num_classes=nc) # 假设有80个类
replace_detect_with_linear(ema.ema, num_classes=nc) # 假设有80个类
# 打印替换后的模型结构
print("Updated model structure after replacing Detect with LinearLayer:")
print(model) # 打印模型结构
# 如果需要,可以打印每一层的详细信息
for name, module in model.named_modules():
print(f"{name}: {module}") # 打印每一层的名称和类型
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
callbacks.run("on_train_epoch_start")
model.train() #
accuracy = 0 # 确保在每个 epoch 开始时初始化准确率
# 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=TQDM_BAR_FORMAT) # progress bar
optimizer.zero_grad()
# print("*"*100)
for i, (imgs, targets, paths, _ , masks, segments) in pbar: # batch -------------------------------------------------------------
callbacks.run("on_train_batch_start") # The name of the supervision is not OK
if opt.test_angle is not None:
if isinstance(opt.test_angle, str):
assert opt.test_angle == "random", f"If test_angle is a string, it must be 'random', got {opt.test_angle}"
random_angle = random.randint(0, 360)
# print(f"targets.shape:{targets.shape}")
# print(f"random_angle:{random_angle}")
# print(f"len(segments):{len(segments)}, len(masks):{len(masks)}, len(imgs):{len(imgs)}, len(targets):{len(targets)}, len(paths):{len(paths)}")
for j in range(len(imgs)):
if len(segments[j]) == 0: #
print(f"j:{j}")
print(f"segments[j] is empty")
print(f"path:{paths[j]}") #
img_np, targets_np, segments_np = apply_test_rotation(im=imgs[j].cpu().permute(1, 2, 0).numpy(), targets=targets[j].cpu().numpy()[1:], segments=segments[j], test_angle=random_angle, return_segments=True)
imgs[j] = torch.from_numpy(img_np).permute(2, 0, 1).to(device) #
# print(f"targets_np.shape:{targets_np.shape}")
# print(targets_np)
targets[j][1:] = torch.from_numpy(targets_np[0]).to(device) #
segments[j] = segments_np #
# function new
# check the rotated bbox
hyper_x = transform_angle(random_angle).to(device) #
else:
assert isinstance(opt.test_angle, (int, float)), f"test_angle must be None, 'random', or a number, got {type(opt.test_angle)}"
if float(opt.test_angle) != 0:
for j in range(len(imgs)):
img_np, targets_np, segments_np = apply_test_rotation(im=imgs[j].cpu().permute(1, 2, 0).numpy(), targets=targets[j].cpu().numpy()[1:], segments=segments[j], test_angle=float(opt.test_angle), return_segments=True)
imgs[j] = torch.from_numpy(img_np).permute(2, 0, 1).to(device)
targets[j][1:] = torch.from_numpy(targets_np[0]).to(device)
segments[j] = segments_np
# ---------------Have a look at one batch of transformed images---------------
# original_imgs = imgs.clone()
# original_targets = targets.clone()
# save_comparison_images(original_imgs, imgs, original_targets, targets, angle, save_dir='tmp_img')
# raise RuntimeError("stop here")
# ---------------Have a look at one batch of transformed images---------------
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(int(imgsz * 0.5), int(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): #
if opt.scn:
# print(hyper_x)
model.model.hyper_forward_and_configure(hyper_x)
pred = model(imgs) # forward
beta1 = model.model.hyper_stack(hyper_x)
angle2 = random.randint(0, 360)
hyper_x2 = transform_angle(angle2).to(device)
beta2 = model.model.hyper_stack(hyper_x2)
else:
# print(type(imgs), imgs.shape) # 确保 imgs 是一个张量
# input()
pred = model(imgs) # forward
# print(f"Output shape of pred: {pred.shape}") # 打印 pred 的形状
# input()
# 输出 targets 的信息
# print(f"Targets shape: {targets.shape}") # 打印 targets 的形状
# print("Targets content:")
# print(targets) # 打印 targets 的内容
# print("------------------------------------------------------------------------")
# 提取第二列作为类别目标
class_targets = targets[:, 1].long() # 提取第二列
# 提取第一列作为排序依据
# sort_indices = targets[:, 0] # 提取第一列
# # 使用第一列的值对第二列进行排序
# sorted_indices = torch.argsort(sort_indices) # 获取排序索引
# sorted_class_targets = class_targets[sorted_indices] # 根据排序索引排序类别目标
# one_hot_targets = torch.nn.functional.one_hot(sorted_class_targets.to(torch.int64), num_classes=nc).float() #
# # if one_hot_targets.dim() == 2: # 如果是 one-hot 编码
# # one_hot_targets = one_hot_targets.argmax(dim=1) # 转换为类别索引
# # 确保 one_hot_targets 的形状与 preds 的输出一致
# if one_hot_targets.dim() == 2: # 如果是 one-hot 编码
# one_hot_targets = one_hot_targets.view(-1, nc) # 确保形状为 [batch_size, num_classes]
# # 打印结果
# print("Sorted class targets:") # 这里是什么内容?
# print(sorted_class_targets) #
# print(f"Sorted class targets shape: {sorted_class_targets.shape}") # 打印大小
# print(f"Sorted class targets one_hot shape: {one_hot_targets.shape}") # 打印大小
# # input() #
# compute_loss now is nn.BCEWithLogitsLoss
# loss = compute_loss(pred, one_hot_targets.to(device)) # loss scaled by batch_size
loss = compute_loss(pred, class_targets.to(device)) # 直接使用类别索引 #
# loss_items = torch.cat((loss, loss, loss)).detach() # test with other stuff
# convert from a scaler to tensor
loss_items = torch.cat((loss.unsqueeze(0), loss.unsqueeze(0), loss.unsqueeze(0))).detach() # 将 loss 转换为一维张量
if opt.scn: #
loss += pow(cos(beta1, beta2),2) #
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.0
# 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
# print(f"The imgsz is {imgsz}")
# input()
results, maps, _ = validate.run_classify(
data_dict,
# batch_size=batch_size // WORLD_SIZE * 2,
batch_size=batch_size,
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,
test_angle=opt.test_angle,
scn=opt.scn,
)
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
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(),
# "opt": vars(opt),
# "git": GIT_INFO, # {remote, branch, commit} if a git repo
# "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_test_angle(value):
try:
return float(value)
except ValueError:
if value == "random":
return value
raise argparse.ArgumentTypeError(f"Invalid value for --test-angle: {value}")
def parse_opt(known=False):
"""
Parse command-line arguments for YOLOv5 training, validation, and testing.
Args:
known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False.
Returns:
(argparse.Namespace): Parsed command-line arguments containing options for YOLOv5 execution.
Example:
```python
from ultralytics.yolo import parse_opt
opt = parse_opt()
print(opt)
```
Links:
- Models: https://github.com/ultralytics/yolov5/tree/master/models
- Datasets: https://github.com/ultralytics/yolov5/tree/master/data
- Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
"""
parser = argparse.ArgumentParser()
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.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=16, 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(
"--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population"
)
parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation")
parser.add_argument("--bucket", type=str, default="", help="gsutil bucket")
parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk")
parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training")
parser.add_argument("--device", default="", 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=8, 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")
# NDJSON logging
parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console")
parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file")
# transformation arguments for SCN
parser.add_argument("--test-angle", type=parse_test_angle, default=None, help="test angle for rotation augmentation. Could be 'random' or a number")
parser.add_argument("--scn-config", type=str, default="config_scn_yolov5.json", help="scn config file path")
parser.add_argument("--scn", action="store_true", help="use scn model")
parser.add_argument("--dimensions", type=int, default=-1, help="Dimensions for SCN model") # 新增的参数
return parser.parse_known_args()[0] if known else parser.parse_args()
def main(opt, callbacks=Callbacks()): #
"""
Runs the main entry point for training or hyperparameter evolution with specified options and optional callbacks.
Args:
opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution.
callbacks (ultralytics.utils.callbacks.Callbacks, optional): Callback functions for various training stages.
Defaults to Callbacks().
Returns:
None
Note:
For detailed usage, refer to:
https://github.com/ultralytics/yolov5/tree/master/models
"""
run = wandb.init(project=f"yolov5_oxford_pet_resort_test_one_layer_classification", name=f"{opt.project}_{opt.dimensions}_{opt.scn_config}", entity="", #
config={"dataset":"oxford_pet"},
)
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
check_requirements(ROOT / "requirements.txt")
# Resume (from specified or most recent last.pt)
if opt.resume 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", timeout=timedelta(seconds=10800) #
)
# Train
if not opt.evolve:
train(opt.hyp, opt, device, callbacks)
# Evolve hyperparameters (optional)
else:
# Hyperparameter evolution metadata (including this hyperparameter True-False, lower_limit, upper_limit)
meta = {
"lr0": (False, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
"lrf": (False, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
"momentum": (False, 0.6, 0.98), # SGD momentum/Adam beta1
"weight_decay": (False, 0.0, 0.001), # optimizer weight decay
"warmup_epochs": (False, 0.0, 5.0), # warmup epochs (fractions ok)
"warmup_momentum": (False, 0.0, 0.95), # warmup initial momentum
"warmup_bias_lr": (False, 0.0, 0.2), # warmup initial bias lr
"box": (False, 0.02, 0.2), # box loss gain
"cls": (False, 0.2, 4.0), # cls loss gain
"cls_pw": (False, 0.5, 2.0), # cls BCELoss positive_weight
"obj": (False, 0.2, 4.0), # obj loss gain (scale with pixels)
"obj_pw": (False, 0.5, 2.0), # obj BCELoss positive_weight
"iou_t": (False, 0.1, 0.7), # IoU training threshold
"anchor_t": (False, 2.0, 8.0), # anchor-multiple threshold
"anchors": (False, 2.0, 10.0), # anchors per output grid (0 to ignore)
"fl_gamma": (False, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
"hsv_h": (True, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
"hsv_s": (True, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
"hsv_v": (True, 0.0, 0.9), # image HSV-Value augmentation (fraction)
"degrees": (True, 0.0, 45.0), # image rotation (+/- deg)
"translate": (True, 0.0, 0.9), # image translation (+/- fraction)
"scale": (True, 0.0, 0.9), # image scale (+/- gain)
"shear": (True, 0.0, 10.0), # image shear (+/- deg)
"perspective": (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
"flipud": (True, 0.0, 1.0), # image flip up-down (probability)
"fliplr": (True, 0.0, 1.0), # image flip left-right (probability)
"mosaic": (True, 0.0, 1.0), # image mosaic (probability)
"mixup": (True, 0.0, 1.0), # image mixup (probability)
"copy_paste": (True, 0.0, 1.0), # segment copy-paste (probability)
}
# GA configs
pop_size = 50
mutation_rate_min = 0.01
mutation_rate_max = 0.5
crossover_rate_min = 0.5
crossover_rate_max = 1
min_elite_size = 2
max_elite_size = 5
tournament_size_min = 2
tournament_size_max = 10
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:
# download evolve.csv if exists
subprocess.run(
[
"gsutil",
"cp",
f"gs://{opt.bucket}/evolve.csv",
str(evolve_csv),
]
)
# Delete the items in meta dictionary whose first value is False
del_ = [item for item, value_ in meta.items() if value_[0] is False]
hyp_GA = hyp.copy() # Make a copy of hyp dictionary
for item in del_:
del meta[item] # Remove the item from meta dictionary
del hyp_GA[item] # Remove the item from hyp_GA dictionary
# Set lower_limit and upper_limit arrays to hold the search space boundaries
lower_limit = np.array([meta[k][1] for k in hyp_GA.keys()])
upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()])
# Create gene_ranges list to hold the range of values for each gene in the population
gene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))]
# Initialize the population with initial_values or random values
initial_values = []
# If resuming evolution from a previous checkpoint
if opt.resume_evolve is not None:
assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!"
with open(ROOT / opt.resume_evolve, errors="ignore") as f:
evolve_population = yaml.safe_load(f)
for value in evolve_population.values():
value = np.array([value[k] for k in hyp_GA.keys()])
initial_values.append(list(value))
# If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_population
else:
yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")]
for file_name in yaml_files:
with open(os.path.join(opt.evolve_population, file_name)) as yaml_file:
value = yaml.safe_load(yaml_file)
value = np.array([value[k] for k in hyp_GA.keys()])
initial_values.append(list(value))
# Generate random values within the search space for the rest of the population
if initial_values is None:
population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)]
elif pop_size > 1:
population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))]
for initial_value in initial_values:
population = [initial_value] + population
# Run the genetic algorithm for a fixed number of generations
list_keys = list(hyp_GA.keys())
for generation in range(opt.evolve):
if generation >= 1:
save_dict = {}
for i in range(len(population)):
little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))}
save_dict[f"gen{str(generation)}number{str(i)}"] = little_dict
with open(save_dir / "evolve_population.yaml", "w") as outfile:
yaml.dump(save_dict, outfile, default_flow_style=False)
# Adaptive elite size
elite_size = min_elite_size + int((max_elite_size - min_elite_size) * (generation / opt.evolve))
# Evaluate the fitness of each individual in the population
fitness_scores = []
for individual in population:
for key, value in zip(hyp_GA.keys(), individual):
hyp_GA[key] = value
hyp.update(hyp_GA)
results = train(hyp.copy(), opt, device, callbacks)
callbacks = Callbacks()
# Write mutation results
keys = (
"metrics/precision",
"metrics/recall",
"metrics/mAP_0.5",
"metrics/mAP_0.5:0.95",
"val/box_loss",
"val/obj_loss",
"val/cls_loss",
)
print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
fitness_scores.append(results[2])
# Select the fittest individuals for reproduction using adaptive tournament selection
selected_indices = []
for _ in range(pop_size - elite_size):
# Adaptive tournament size
tournament_size = max(
max(2, tournament_size_min),
int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))),
)
# Perform tournament selection to choose the best individual
tournament_indices = random.sample(range(pop_size), tournament_size)
tournament_fitness = [fitness_scores[j] for j in tournament_indices]
winner_index = tournament_indices[tournament_fitness.index(max(tournament_fitness))]
selected_indices.append(winner_index)
# Add the elite individuals to the selected indices
elite_indices = [i for i in range(pop_size) if fitness_scores[i] in sorted(fitness_scores)[-elite_size:]]
selected_indices.extend(elite_indices)
# Create the next generation through crossover and mutation
next_generation = []
for _ in range(pop_size):
parent1_index = selected_indices[random.randint(0, pop_size - 1)]
parent2_index = selected_indices[random.randint(0, pop_size - 1)]
# Adaptive crossover rate
crossover_rate = max(
crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve))
)
if random.uniform(0, 1) < crossover_rate:
crossover_point = random.randint(1, len(hyp_GA) - 1)
child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:]
else:
child = population[parent1_index]
# Adaptive mutation rate
mutation_rate = max(
mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve))
)
for j in range(len(hyp_GA)):
if random.uniform(0, 1) < mutation_rate:
child[j] += random.uniform(-0.1, 0.1)
child[j] = min(max(child[j], gene_ranges[j][0]), gene_ranges[j][1])
next_generation.append(child)
# Replace the old population with the new generation
population = next_generation
# Print the best solution found
best_index = fitness_scores.index(max(fitness_scores))
best_individual = population[best_index]
print("Best solution found:", best_individual)
# 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 generate_individual(input_ranges, individual_length):
"""
Generate an individual with random hyperparameters within specified ranges.
Args:
input_ranges (list[tuple[float, float]]): List of tuples where each tuple contains the lower and upper bounds
for the corresponding gene (hyperparameter).
individual_length (int): The number of genes (hyperparameters) in the individual.
Returns:
list[float]: A list representing a generated individual with random gene values within the specified ranges.
Example:
```python
input_ranges = [(0.01, 0.1), (0.1, 1.0), (0.9, 2.0)]
individual_length = 3
individual = generate_individual(input_ranges, individual_length)
print(individual) # Output: [0.035, 0.678, 1.456] (example output)
```
Note:
The individual returned will have a length equal to `individual_length`, with each gene value being a floating-point
number within its specified range in `input_ranges`.
"""
individual = []
for i in range(individual_length):
lower_bound, upper_bound = input_ranges[i]
individual.append(random.uniform(lower_bound, upper_bound))
return individual
def run(**kwargs):
"""
Execute YOLOv5 training with specified options, allowing optional overrides through keyword arguments.
Args:
weights (str, optional): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'.
cfg (str, optional): Path to model YAML configuration. Defaults to an empty string.
data (str, optional): Path to dataset YAML configuration. Defaults to ROOT / 'data/coco128.yaml'.
hyp (str, optional): Path to hyperparameters YAML configuration. Defaults to ROOT / 'data/hyps/hyp.scratch-low.yaml'.
epochs (int, optional): Total number of training epochs. Defaults to 100.
batch_size (int, optional): Total batch size for all GPUs. Use -1 for automatic batch size determination. Defaults to 16.
imgsz (int, optional): Image size (pixels) for training and validation. Defaults to 640.
rect (bool, optional): Use rectangular training. Defaults to False.
resume (bool | str, optional): Resume most recent training with an optional path. Defaults to False.
nosave (bool, optional): Only save the final checkpoint. Defaults to False.
noval (bool, optional): Only validate at the final epoch. Defaults to False.
noautoanchor (bool, optional): Disable AutoAnchor. Defaults to False.
noplots (bool, optional): Do not save plot files. Defaults to False.
evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a
value.
evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/ hyps'.
resume_evolve (str, optional): Resume hyperparameter evolution from the last generation. Defaults to None.
bucket (str, optional): gsutil bucket for saving checkpoints. Defaults to an empty string.
cache (str, optional): Cache image data in 'ram' or 'disk'. Defaults to None.
image_weights (bool, optional): Use weighted image selection for training. Defaults to False.
device (str, optional): CUDA device identifier, e.g., '0', '0,1,2,3', or 'cpu'. Defaults to an empty string.
multi_scale (bool, optional): Use multi-scale training, varying image size by ±50%. Defaults to False.
single_cls (bool, optional): Train with multi-class data as single-class. Defaults to False.
optimizer (str, optional): Optimizer type, choices are ['SGD', 'Adam', 'AdamW']. Defaults to 'SGD'.
sync_bn (bool, optional): Use synchronized BatchNorm, only available in DDP mode. Defaults to False.
workers (int, optional): Maximum dataloader workers per rank in DDP mode. Defaults to 8.
project (str, optional): Directory for saving training runs. Defaults to ROOT / 'runs/train'.
name (str, optional): Name for saving the training run. Defaults to 'exp'.
exist_ok (bool, optional): Allow existing project/name without incrementing. Defaults to False.
quad (bool, optional): Use quad dataloader. Defaults to False.
cos_lr (bool, optional): Use cosine learning rate scheduler. Defaults to False.
label_smoothing (float, optional): Label smoothing epsilon value. Defaults to 0.0.
patience (int, optional): Patience for early stopping, measured in epochs without improvement. Defaults to 100.
freeze (list, optional): Layers to freeze, e.g., backbone=10, first 3 layers = [0, 1, 2]. Defaults to [0].
save_period (int, optional): Frequency in epochs to save checkpoints. Disabled if < 1. Defaults to -1.
seed (int, optional): Global training random seed. Defaults to 0.
local_rank (int, optional): Automatic DDP Multi-GPU argument. Do not modify. Defaults to -1.
Returns:
None: The function initiates YOLOv5 training or hyperparameter evolution based on the provided options.
Examples:
```python
import train
train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
```
Notes:
- Models: https://github.com/ultralytics/yolov5/tree/master/models
- Datasets: https://github.com/ultralytics/yolov5/tree/master/data
- Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
"""
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)
# this is good
# CUDA_VISIBLE_DEVICES=3 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 150 --hyp data/hyps/hyp.scratch-low-stanford-dog.yaml --cache ram --optimizer Adam --workers 6 --batch-size 128 --patience 50 --device 3
# with cos-lr
# CUDA_VISIBLE_DEVICES=7 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 300 --hyp data/hyps/hyp.scratch-low-stanford-dog.yaml --cache ram --optimizer Adam --workers 6 --batch-size 128 --cos-lr --device 7
# with patience
# CUDA_VISIBLE_DEVICES=7 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 100 --patience 5 --hyp data/hyps/hyp.scratch-low-stanford-dog.yaml --cache ram --optimizer Adam --workers 12 --batch-size 128 --device 7
# CUDA_VISIBLE_DEVICES=7 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 100 --patience 10 --hyp data/hyps/hyp.scratch-low-stanford-dog.yaml --cache ram --optimizer Adam --workers 12 --batch-size 128 --device 7 --test-angle random
# CUDA_VISIBLE_DEVICES=7 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 100 --patience 5 --hyp data/hyps/hyp.scratch-low-stanford-dog.yaml --cache ram --optimizer Adam --workers 12 --batch-size 128 --device 7 --scn
# naive
# CUDA_VISIBLE_DEVICES=3 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 100 --cache ram --optimizer Adam --workers 6 --batch-size 128 --device 3
# evolve for cos-lr
# CUDA_VISIBLE_DEVICES=3 python train.py --project naive --data data/stanford_dogs.yaml --cfg yolov5s.yaml --weights '' --img 320 --epochs 10 --cache ram --optimizer Adam --workers 6 --batch-size 128 --cos-lr --device 3 --evolve
CUDA_VISIBLE_DEVICES=0 python train_yudan_new_setting_scn_one_path_classification.py --project test_deep_layer_one4one --data data/oxford_pet_yudan.yaml --cfg yolov5s_deep_level_scale.yaml --weights weights/yolov5s.pt --img 320 --epochs 100 --patience 1000 --cache ram --optimizer Adam --workers 12 --batch-size 128 --hyp data/hyps/hyp.scratch-low-oxford-pet1.yaml --device 0
@smart_inference_mode()
def run_classify( # the real name is % and OK
data,
weights=None, # model.pt path(s)
batch_size=32, # batch size
imgsz=640, # inference size (pixels)
conf_thres=0.001, # confidence threshold
iou_thres=0.6, # NMS IoU threshold
max_det=300, # maximum detections per image
task="val", # train, val, test, speed or study
device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu
workers=8, # max dataloader workers (per RANK in DDP mode)
single_cls=False, # treat as single-class dataset
augment=False, # augmented inference
verbose=False, # verbose output
save_txt=False, # save results to *.txt
save_hybrid=False, # save label+prediction hybrid results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_json=False, # save a COCO-JSON results file
project=ROOT / "runs/val", # save to project/name
name="exp", # save to project/name
exist_ok=False, # existing project/name ok, do not increment
half=True, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
model=None,
dataloader=None,
save_dir=Path(""),
plots=True,
callbacks=Callbacks(),
compute_loss=None,
test_angle=None,
scn=False,
):
"""
Evaluates a YOLOv5 model on a dataset and logs performance metrics.
Args:
data (str | dict): Path to a dataset YAML file or a dataset dictionary.
weights (str | list[str], optional): Path to the model weights file(s). Supports various formats including PyTorch,
TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite,
TensorFlow Edge TPU, and PaddlePaddle.
batch_size (int, optional): Batch size for inference. Default is 32.
imgsz (int, optional): Input image size (pixels). Default is 640.
conf_thres (float, optional): Confidence threshold for object detection. Default is 0.001.
iou_thres (float, optional): IoU threshold for Non-Maximum Suppression (NMS). Default is 0.6.
max_det (int, optional): Maximum number of detections per image. Default is 300.
task (str, optional): Task type - 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'.
device (str, optional): Device to use for computation, e.g., '0' or '0,1,2,3' for CUDA or 'cpu' for CPU. Default is ''.
workers (int, optional): Number of dataloader workers. Default is 8.
single_cls (bool, optional): Treat dataset as a single class. Default is False.
augment (bool, optional): Enable augmented inference. Default is False.
verbose (bool, optional): Enable verbose output. Default is False.
save_txt (bool, optional): Save results to *.txt files. Default is False.
save_hybrid (bool, optional): Save label and prediction hybrid results to *.txt files. Default is False.
save_conf (bool, optional): Save confidences in --save-txt labels. Default is False.
save_json (bool, optional): Save a COCO-JSON results file. Default is False.
project (str | Path, optional): Directory to save results. Default is ROOT/'runs/val'.
name (str, optional): Name of the run. Default is 'exp'.
exist_ok (bool, optional): Overwrite existing project/name without incrementing. Default is False.
half (bool, optional): Use FP16 half-precision inference. Default is True.
dnn (bool, optional): Use OpenCV DNN for ONNX inference. Default is False.
model (torch.nn.Module, optional): Model object for training. Default is None.
dataloader (torch.utils.data.DataLoader, optional): Dataloader object. Default is None.
save_dir (Path, optional): Directory to save results. Default is Path('').
plots (bool, optional): Plot validation images and metrics. Default is True.
callbacks (utils.callbacks.Callbacks, optional): Callbacks for logging and monitoring. Default is Callbacks().
compute_loss (function, optional): Loss function for training. Default is None.
Returns:
dict: Contains performance metrics including precision, recall, mAP50, and mAP50-95.
"""
# Initialize/load model and set device
training = model is not None #
if training: # called by train.py
device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model
half &= device.type != "cpu" # half precision only supported on CUDA
model.half() if half else model.float()
else: # called directly
device = select_device(device, batch_size=batch_size) #
# Directories
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
(save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Load model
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
imgsz = check_img_size(imgsz, s=stride) # check image size
# print(f"The size of the imgsz is {imgsz}")
# input()
half = model.fp16 # FP16 supported on limited backends with CUDA
if engine:
batch_size = model.batch_size
else:
device = model.device
if not (pt or jit):
batch_size = 1 # export.py models default to batch-size 1
LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models")
# Data
data = check_dataset(data) # check
# Configure
model.eval()
cuda = device.type != "cpu"
is_coco = isinstance(data.get("val"), str) and data["val"].endswith(f"coco{os.sep}val2017.txt") # COCO dataset
nc = 1 if single_cls else int(data["nc"]) # number of classes
iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training: #
if pt and not single_cls: # check --weights are trained on --data
ncm = model.model.nc
assert ncm == nc, (
f"{weights} ({ncm} classes) trained on different --data than what you passed ({nc} "
f"classes). Pass correct combination of --weights and --data that are trained together."
)
model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz)) # warmup
pad, rect = (0.0, False) if task == "speed" else (0.5, pt) # square inference for benchmarks
task = task if task in ("train", "val", "test") else "val" # path to train/val/test images
dataloader = create_dataloader( #
data[task],
imgsz,
batch_size,
stride,
single_cls,
pad=pad,
rect=rect,
workers=workers,
prefix=colorstr(f"{task}: "),
)[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = model.names if hasattr(model, "names") else model.module.names # get class names
if isinstance(names, (list, tuple)): # old format
names = dict(enumerate(names))
class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
s = ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "P", "R", "mAP50", "mAP50-95")
tp, fp, p, r, f1, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
dt = Profile(device=device), Profile(device=device), Profile(device=device) # profiling times
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
callbacks.run("on_val_start")
test_loss, accuracy = 0, 0
pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT) # progress bar
for batch_i, (im, targets, paths, shapes, masks, segments) in enumerate(pbar):
callbacks.run("on_val_batch_start")
if test_angle is not None:
if isinstance(test_angle, str):
assert test_angle == "random", f"If test_angle is a string, it must be 'random', got {test_angle}"
# apply random rotation augmentation
random_angle = random.randint(0, 360)
# im, targets, paths = apply_rotation_augmentation(im, targets, paths, random_angle)
for j in range(len(im)):
img_np, targets_np, segments_np = apply_test_rotation(im=im[j].cpu().permute(1, 2, 0).numpy(), targets=targets[j].cpu().numpy()[1:], segments=segments[j], test_angle=random_angle, return_segments=True)
im[j] = torch.from_numpy(img_np).permute(2, 0, 1).to(device)
targets[j][1:] = torch.from_numpy(targets_np[0]).to(device)
segments[j] = segments_np
if scn:
hyper_x = transform_angle(random_angle).to(device)
else:
assert isinstance(test_angle, (int, float)), f"test_angle must be None, 'random', or a number, got {type(test_angle)}"
# apply rotation augmentation
# im, targets, paths = apply_rotation_augmentation(im, targets, paths, test_angle)
if float(test_angle) != 0:
for j in range(len(im)):
# test if we are in this branch
input() #
img_np, targets_np, segments_np = apply_test_rotation(im=im[j].cpu().permute(1, 2, 0).numpy(), targets=targets[j].cpu().numpy()[1:], segments=segments[j], test_angle=float(test_angle), return_segments=True)
im[j] = torch.from_numpy(img_np).permute(2, 0, 1).to(device)
targets[j][1:] = torch.from_numpy(targets_np[0]).to(device)
segments[j] = segments_np
with dt[0]: #
if cuda:
im = im.to(device, non_blocking=True)
targets = targets.to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
nb, _, height, width = im.shape # batch size, channels, height, width
# Inference
with dt[1]:
if scn:
model.model.hyper_forward_and_configure(hyper_x.half())
# preds, train_out = model(im) if compute_loss else (model(im, augment=augment), None)
# 输出 im 和 preds 的大小
# print(f"Input image size (im): {im.size()}") # 输出 im 的大小
preds = model(im) # only do the compute loss when training
# print(f"Predictions size (preds): {preds.size()}") # 输出 preds 的大小
# input()
# Loss
if compute_loss:
class_targets = targets[:, 1].long() # 提取第二列
# 提取第一列作为排序依据
# sort_indices = targets[:, 0] # 提取第一列
# # 使用第一列的值对第二列进行排序
# sorted_indices = torch.argsort(sort_indices) # 获取排序索引
# sorted_class_targets = class_targets[sorted_indices] # 根据排序索引排序类别目标
# one_hot_targets = torch.nn.functional.one_hot(sorted_class_targets.to(torch.int64), num_classes=nc).float()
# # 确保 one_hot_targets 的形状与 preds 的输出一致
# if one_hot_targets.dim() == 2: # 如果是 one-hot 编码
# one_hot_targets = one_hot_targets.argmax(dim=1) # 转换为类别索引
loss += compute_loss(preds, class_targets) # box, obj, cls
# 计算准确率
# accuracy += (preds.argmax(1) == one_hot_targets).type(torch.float).sum().item() # the acc one # 感觉有点小问题,很奇怪
accuracy += (preds.argmax(1) == class_targets).type(torch.float).sum().item() # the acc one # 感觉有点小问题,很奇怪
# print(f"The accuracy is {accuracy}")
# input()
# Metrics
# 这里还需要进行修改
for si, pred in enumerate(preds): #
labels = targets[targets[:, 0] == si, 1:] # 获取当前样本的标签
nl, npr = labels.shape[0], pred.shape[0] # number of labels, predictions
path, shape = Path(paths[si]), shapes[si][0]
seen += 1
if npr == 0:
if nl:
stats.append((torch.zeros(npr, niou, dtype=torch.bool, device=device), *torch.zeros((2, 0), device=device), labels[:, 0]))
if plots:
confusion_matrix.process_batch(detections=None, labels=labels[:, 0]) #
continue
# Predictions
# pred_probs = pred.softmax(dim=1) # 对于多类分类,使用 softmax 获取概率 #
# pred_classes = pred.argmax(1) # 获取每个样本的预测类别
# 计算准确率
# accuracy = (pred_classes == labels[:, 0].long()).float().mean().item()
# accuracy += (pred.argmax(1) == one_hot_targets).type(torch.float).sum().item()
# 计算精确率
# true_positives = (pred_classes == 1) & (labels[:, 0].long() == 1) # 假设正类为 1
# predicted_positives = pred_classes == 1
precision = 0 # 加上小常数以避免除零
# 打印或记录准确率和精确率
# print(f"Sample {si}: Accuracy: {accuracy:.4f}, Precision: {precision:.4f}")
# Evaluate
# if nl:
# 这里不再需要处理目标框,因为我们只关注分类
# correct = (pred_classes == labels[:, 0].long()).float() # 计算正确预测
# stats.append((correct, pred[:, 4], pred[:, 5], labels[:, 0])) # 这里可以根据需要调整
# Save/log
if save_txt:
(save_dir / "labels").mkdir(parents=True, exist_ok=True)
save_one_txt(pred, save_conf, shape, file=save_dir / "labels" / f"{path.stem}.txt")
if save_json:
save_one_json(pred, jdict, path, class_map) # append to COCO-JSON dictionary
callbacks.run("on_val_image_end", pred, pred, path, names, im[si])
# print(f"The acc is {accuracy}") #
acc_percent = accuracy/(len(dataloader.dataset))
# print()
print(f"Accuracy: {(100*acc_percent):>0.1f}%")
print(f"Loss: {loss}")
# input()
# Compute metrics
stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
acc, tp, fp, p, r, f1, ap, ap_class = accuracy, ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names) #
r = acc # use the recall to represenet the accuracy #
ap50, ap = ap[:, 0], ap.mean(1) # [email protected], [email protected]:0.95
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
# nt = np.bincount(stats[3].astype(int), minlength=nc) # number of targets per class #
nt = 1 # hard set
# Print results
pf = "%22s" + "%11i" * 2 + "%11.3g" * 4 # print format
LOGGER.info(pf % ("all", seen, nt, mp, mr, map50, map)) # 直接使用 nt
if nt == 0:
LOGGER.warning(f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels")
# Print speeds
t = tuple(x.t / seen * 1e3 for x in dt) # speeds per image
if not training:
shape = (batch_size, 3, imgsz, imgsz) #
LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}" % t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
callbacks.run("on_val_end", nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else "" # weights
anno_json = str(Path("../datasets/coco/annotations/instances_val2017.json")) # annotations
if not os.path.exists(anno_json):
anno_json = os.path.join(data["path"], "annotations", "instances_val2017.json")
pred_json = str(save_dir / f"{w}_predictions.json") # predictions
LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...")
with open(pred_json, "w") as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
check_requirements("pycocotools>=2.0.6")
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, "bbox")
if is_coco:
eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
LOGGER.info(f"pycocotools unable to run: {e}")
# Return results
model.float() # for training
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment