前言
源码: YOLOv5源码. 导航: 【YOLOV5-5.0 源码讲解】整体项目文件导航.
这个文件主要是在每一轮训练结束后,验证当前模型的mAP、混淆矩阵等指标。这个文件本来是叫test.py的,但我觉得叫val.py可能更加恰当些,所以改了。
实际上这个脚本最常用的应该是通过train.py调用 run 函数,而不是通过执行 val.py 的。所以在了解这个脚本的时候,其实最重要的就是 run 函数。
0、导入需要的包和基本配置
import argparse
import json
import os
import sys
from pathlib import Path
from threading import Thread
import numpy as np
import torch
import yaml
from tqdm import tqdm
FILE = Path(__file__).absolute()
sys.path.append(FILE.parents[0].as_posix())
from models.experimental import attempt_load
from utils.datasets import create_dataloader
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \
box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
from utils.metrics import ap_per_class, ConfusionMatrix
from utils.plots import plot_images, output_to_target, plot_study_txt
from utils.torch_utils import select_device, time_synchronized
1、设置opt参数
def parse_opt():
"""
opt参数详解
data: 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息
weights: 模型的权重文件地址 weights/yolov5s.pt
batch_size: 前向传播的批次大小 默认32
imgsz: 输入网络的图片分辨率 默认640
conf-thres: object置信度阈值 默认0.25
iou-thres: 进行NMS时IOU的阈值 默认0.6
task: 设置测试的类型 有train, val, test, speed or study几种 默认val
device: 测试的设备
single-cls: 数据集是否只用一个类别 默认False
augment: 测试是否使用TTA Test Time Augment 默认False
verbose: 是否打印出每个类别的mAP 默认False
save-txt: 是否以txt文件的形式保存模型预测框的坐标 默认True
save-hybrid: 是否save label+prediction hybrid results to *.txt 默认False
save-conf: 是否保存预测每个目标的置信度到预测tx文件中 默认True
save-json: 是否按照coco的json格式保存预测框,并且使用cocoapi做评估(需要同样coco的json格式的标签) 默认False
project: 测试保存的源文件 默认runs/test
name: 测试保存的文件地址 默认exp 保存在runs/test/exp下
exist-ok: 是否存在当前文件 默认False 一般是 no exist-ok 连用 所以一般都要重新创建文件夹
half: 是否使用半精度推理 默认False
"""
parser = argparse.ArgumentParser(prog='val.py')
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='dataset.yaml path')
parser.add_argument('--weights', nargs='+', type=str, default='weights/yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--batch-size', type=int, default=4, help='batch size')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
parser.add_argument('--task', default='val', help='train, val, test, speed or study')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--verbose', action='store_true', help='report mAP by class')
parser.add_argument('--save-txt', default=True, action='store_true', help='save results to *.txt')
parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
parser.add_argument('--save-conf', default=True, action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
parser.add_argument('--project', default='runs/test', 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('--half', action='store_true', help='use FP16 half-precision inference')
opt = parser.parse_args()
opt.save_json |= opt.data.endswith('coco.yaml')
opt.save_txt |= opt.save_hybrid
opt.data = check_file(opt.data)
return opt
2、执行main函数
def main(opt):
set_logging()
print(colorstr('test: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
check_requirements(exclude=('tensorboard', 'thop'))
if opt.task in ('train', 'val', 'test'):
run(**vars(opt))
elif opt.task == 'speed':
for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
run(opt.data, weights=w, batch_size=opt.batch_size, imgsz=opt.imgsz, conf_thres=.25, iou_thres=.45,
save_json=False, plots=False)
elif opt.task == 'study':
x = list(range(256, 1536 + 128, 128))
for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt'
y = []
for i in x:
print(f'\nRunning {f} point {i}...')
r, _, t = run(opt.data, weights=w, batch_size=opt.batch_size, imgsz=i, conf_thres=opt.conf_thres,
iou_thres=opt.iou_thres, save_json=opt.save_json, plots=False)
y.append(r + t)
np.savetxt(f, y, fmt='%10.4g')
os.system('zip -r study.zip study_*.txt')
plot_study_txt(x=x)
可以看到这个模块根据opt.task可以分为三个分支,我们主要的分支还是在 opt.task in (‘train’, ‘val’, ‘test’)。而其他的两个分支,大家大概看看在干什么就可以了,没什么用。一般我们都是直接进入第一个分支,执行run函数。
3、执行run函数
run 函数其实用train.py执行的,并不是执行val.py。
train.py调用(每个训练epoch后验证当前模型):
3.1、载入参数
@torch.no_grad()
def run(data, weights=None, batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.6,
task='val', device='', single_cls=False, augment=False, verbose=False, save_txt=False,
save_hybrid=False, save_conf=False, save_json=False, project='runs/test', name='exp',
exist_ok=False, half=True, model=None, dataloader=None, save_dir=Path(''), plots=True,
wandb_logger=None, compute_loss=None,
):
"""
:params data: 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息 train.py时传入data_dict
:params weights: 模型的权重文件地址 运行train.py=None 运行test.py=默认weights/yolov5s.pt
:params batch_size: 前向传播的批次大小 运行test.py传入默认32 运行train.py则传入batch_size // WORLD_SIZE * 2
:params imgsz: 输入网络的图片分辨率 运行test.py传入默认640 运行train.py则传入imgsz_test
:params conf_thres: object置信度阈值 默认0.25
:params iou_thres: 进行NMS时IOU的阈值 默认0.6
:params task: 设置测试的类型 有train, val, test, speed or study几种 默认val
:params device: 测试的设备
:params single_cls: 数据集是否只用一个类别 运行test.py传入默认False 运行train.py则传入single_cls
:params augment: 测试是否使用TTA Test Time Augment 默认False
:params verbose: 是否打印出每个类别的mAP 运行test.py传入默认Fasle 运行train.py则传入nc < 50 and final_epoch
:params save_txt: 是否以txt文件的形式保存模型预测框的坐标 默认True
:params save_hybrid: 是否save label+prediction hybrid results to *.txt 默认False
:params save_conf: 是否保存预测每个目标的置信度到预测tx文件中 默认True
:params save_json: 是否按照coco的json格式保存预测框,并且使用cocoapi做评估(需要同样coco的json格式的标签)
运行test.py传入默认Fasle 运行train.py则传入is_coco and final_epoch(一般也是False)
:params project: 测试保存的源文件 默认runs/test
:params name: 测试保存的文件地址 默认exp 保存在runs/test/exp下
:params exist_ok: 是否存在当前文件 默认False 一般是 no exist-ok 连用 所以一般都要重新创建文件夹
:params half: 是否使用半精度推理 FP16 half-precision inference 默认False
:params model: 模型 如果执行test.py就为None 如果执行train.py就会传入ema.ema(ema模型)
:params dataloader: 数据加载器 如果执行test.py就为None 如果执行train.py就会传入testloader
:params save_dir: 文件保存路径 如果执行test.py就为‘’ 如果执行train.py就会传入save_dir(runs/train/expn)
:params plots: 是否可视化 运行test.py传入默认True 运行train.py则传入plots and final_epoch
:params wandb_logger: 网页可视化 类似于tensorboard 运行test.py传入默认None 运行train.py则传入wandb_logger(train)
:params compute_loss: 损失函数 运行test.py传入默认None 运行train.py则传入compute_loss(train)
:return (Precision, Recall, map@0.5, map@0.5:0.95, box_loss, obj_loss, cls_loss)
"""
3.2、初始化配置1
训练时(train.py)调用:初始化模型参数、训练设备 验证时(val.py)调用:初始化设备、save_dir文件路径、make dir、加载模型、check imgsz、 加载+check data配置信息
training = model is not None
if training:
device = next(model.parameters()).device
else:
device = select_device(device, batch_size=batch_size)
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)
model = attempt_load(weights, map_location=device)
gs = max(int(model.stride.max()), 32)
imgsz = check_img_size(imgsz, s=gs)
with open(data) as f:
data = yaml.safe_load(f)
check_dataset(data)
3.3、调整模型
半精度验证half model + 模型剪枝prune + 模型融合conv+bn
half &= device.type != 'cpu'
if half:
model.half()
model.eval()
3.4、初始化配置2
是否是coco数据集is_coco + 类别个数nc + 计算mAP相关参数 + 初始化日志 Logging
is_coco = type(data['val']) is str and data['val'].endswith('coco/val2017.txt')
nc = 1 if single_cls else int(data['nc'])
iouv = torch.linspace(0.5, 0.95, 10).to(device)
niou = iouv.numel()
log_imgs = 0
if wandb_logger and wandb_logger.wandb:
log_imgs = min(wandb_logger.log_imgs, 100)
3.5、加载val数据集
训练时(train.py)调用:加载val数据集 验证时(val.py)调用:不需要加载val数据集 直接从train.py 中传入testloader
if not training:
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
task = task if task in ('train', 'val', 'test') else 'val'
dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=0.5, rect=True,
prefix=colorstr(f'{task}: '))[0]
3.6、初始化配置3
初始化混淆矩阵 + 数据集类名 + 获取coco数据集的类别索引 + 设置tqdm进度条 + 初始化p, r, f1, mp, mr, map50, map指标和时间t0, t1, t2 + 初始化测试集的损失 + 初始化json文件中的字典 统计信息 ap等
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
p, r, f1, mp, mr, map50, map, t0, t1, t2 = 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
3.7、开始验证
for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
3.7.1、预处理图片和target
t_ = time_synchronized()
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float()
img /= 255.0
targets = targets.to(device)
nb, _, height, width = img.shape
t = time_synchronized()
t0 += t - t_
3.7.2、model 前向推理
out, train_out = model(img, augment=augment)
t1 += time_synchronized() - t
3.7.3、计算验证集损失
if compute_loss:
loss += compute_loss([x.float() for x in train_out], targets)[1][:3]
3.7.4、Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []
t = time_synchronized()
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
t2 += time_synchronized() - t
3.7.5、统计每张图片的真实框、预测框信息
for si, pred in enumerate(out):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else []
path = Path(paths[si])
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
continue
if single_cls:
pred[:, 5] = 0
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1])
3.7.6、保存预测信息到image_name.txt文件
预测信息:cls, xywh, conf
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]]
for *xyxy, conf, cls in predn.tolist():
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
line = (cls, *xywh, conf) if save_conf else (cls, *xywh)
with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
如下图:
3.7.7、保存预测信息到wandb_logger
if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0:
if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": int(cls),
"box_caption": "%s %.3f" % (names[cls], conf),
"scores": {"class_score": conf},
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
boxes = {"predictions": {"box_data": box_data, "class_labels": names}}
wandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))
wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None
3.7.8、将预测信息保存到coco格式的json字典
jdict字典(预测信息):image_id + category_id + bbox + score
if save_json:
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4])
box[:, :2] -= box[:, 2:] / 2
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({'image_id': image_id,
'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})
3.7.9、计算混淆矩阵、计算correct、生成stats
correct(TP)计算过程
for gt中所有类别:
- 选出pred中属于该类别的所有预测框
- 选出gt中属于该类别的所有gt框
- 计算出选出的所有预测框 和 选出的所有gt框 ious
- 筛选出所有ious > 0.5的预测框 就是TP
- 如果存在TP 就统计所有TP中不同iou阈值下的TP 同时统计检测到的目标(detected)
- 重复这个过程 直到检测到的目标个数len(detected) = gt个数
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
if nl:
detected = []
tcls_tensor = labels[:, 0]
tbox = xywh2xyxy(labels[:, 1:5])
scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1])
if plots:
confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1)
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1)
if pi.shape[0]:
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1)
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]]
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
correct[pi[j]] = ious[j] > iouv
if len(detected) == nl:
break
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
3.7.10、画出前三个batch图片的gt和pred框
if plots and batch_i < 3:
f = save_dir / f'test_batch{batch_i}_labels.jpg'
Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()
f = save_dir / f'test_batch{batch_i}_pred.jpg'
Thread(target=plot_images, args=(img, output_to_target(out), paths, f, names), daemon=True).start()
3.7.11、计算mAP
stats = [np.concatenate(x, 0) for x in zip(*stats)]
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
ap50, ap = ap[:, 0], ap.mean(1)
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64), minlength=nc)
else:
nt = torch.zeros(1)
3.7.12、print打印各项指标
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
t = tuple(x / seen * 1E3 for x in (t0, t1, t2))
if not training:
shape = (batch_size, 3, imgsz, imgsz)
print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
3.7.13、画出混淆矩阵并存入wandb_logger中
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
if wandb_logger and wandb_logger.wandb:
val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
wandb_logger.log({"Validation": val_batches})
if wandb_images:
wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})
3.7.14、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 ''
anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json')
pred_json = str(save_dir / f"{w}_predictions.json")
print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try:
check_requirements(['pycocotools'])
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json)
pred = anno.loadRes(pred_json)
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files]
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:2]
except Exception as e:
print(f'pycocotools unable to run: {e}')
3.7.15、 Return results
model.float()
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {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
|