Created
January 23, 2025 13:44
-
-
Save dyd1234/f4faad5a5fbf35aaccb1d3bb1a8b630f to your computer and use it in GitHub Desktop.
Train for YOLO v5s with convering the detection head to linear layer
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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) | |
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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 | |
# 创建 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
@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