Merge pull request #149 from Jiang-Stan/jst/add_yolov5_example

Jst/add yolov5 example

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "examples/quantization_aware_training/imagenet1k/deit/deit"]
  path = examples/quantization_aware_training/imagenet1k/deit/deit
  url = https://github.com/facebookresearch/deit.git
+[submodule "examples/post_training_quantization/coco2017/yolov5/yolov5"]
+ path = examples/post_training_quantization/coco2017/yolov5/yolov5
+ url = https://github.com/ultralytics/yolov5.git

examples/post_training_quantization/coco2017/yolov5/README.md

@@ -0,0 +1,63 @@
+## Command
+```
+python3 main.py --model_name yolov5n --qconfig_path qconfig.yaml --data_path /PATH/TO/COCO --checkpoint_path checkpoints/yolov5n.pth
+python3 main.py --model_name yolov5s --qconfig_path qconfig.yaml --data_path /PATH/TO/COCO --checkpoint_path checkpoints/yolov5s.pth
+```
+
+## Dataset
+- Download and prepare coco2017 as described in YOLOv5 repo, which should have the following basic structure:
+
+ ```
+ coco
+ └── images
+ └── train2017
+ └── val2017
+ └── annotations
+ └── labels
+ └── train2017.txt
+ └── val2017.txt
+ ```
+
+## Calibration
+- Random sample image paths for calibration:
+
+ ```
+ python3 random_sample_calib.py --data_path /PATH/TO/COCO
+ ```
+
+## Pretraind model
+- create checkpoints dir:
+ ```
+ mkdir ./checkpoints
+ ```
+- Download float checkpoints:
+ - [yolov5n](https://drive.google.com/file/d/1pcsVQHoHCZ4N0ZB8E2QfDFzCmKfSCOjz/view?usp=sharing)
+ - [yolov5s](https://drive.google.com/file/d/1fsDtQtnmNfMM6n0CpslzTMca7xkiaWhq/view?usp=sharing)
+
+## Requirements
+```
+pip install -r yolov5/requirements.txt
+```
+
+## COCO Benchmark
+- Task: COCO
+- Eval data num: 5k
+- Calibration data num: 128
+- Weight bit: 8
+- Feature bit: 8
+- Weight
+ - Granularity: channel-wise
+ - Scheme: symmetric
+ - Observer: MinMax
+- Feature
+ - Granularity: tensor-wise
+ - Scheme: asymmetric
+ - Observer: MinMax
+
+|Model|qconfig|mAP50-95|mAP50|prec|recall|
+|-----|-----|-----|-----|-----|-----|
+|YOLOv5n|float|27.7%|45.6%|57.5%|43.2%|
+|YOLOv5n|8w8f|27.3%|45.2%|58.0%|42.8%|
+||
+|YOLOv5s|float|37.1%|56.6%|66.8%|52.1%|
+|YOLOv5s|8w8f|36.7%|56.5%|66.2%|52.1%|

examples/post_training_quantization/coco2017/yolov5/main.py

@@ -0,0 +1,303 @@
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+"""
+Validate a trained YOLOv5 detection model on a detection dataset
+
+Usage:
+ $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
+
+Usage - formats:
+ $ python val.py --weights yolov5s.pt # PyTorch
+ yolov5s.torchscript # TorchScript
+ yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
+ yolov5s_openvino_model # OpenVINO
+ yolov5s.engine # TensorRT
+ yolov5s.mlmodel # CoreML (macOS-only)
+ yolov5s_saved_model # TensorFlow SavedModel
+ yolov5s.pb # TensorFlow GraphDef
+ yolov5s.tflite # TensorFlow Lite
+ yolov5s_edgetpu.tflite # TensorFlow Edge TPU
+ yolov5s_paddle_model # PaddlePaddle
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+FILE = Path(__file__).resolve()
+ROOT = str(FILE.parents[0]) + "/yolov5" # YOLOv5 root directory
+if ROOT not in sys.path:
+ sys.path.append(ROOT) # add ROOT to PATH
+
+from models import yolov5n, yolov5s
+from yolov5.utils.dataloaders import create_dataloader
+from yolov5.utils.general import (
+ LOGGER,
+ TQDM_BAR_FORMAT,
+ Profile,
+ colorstr,
+ non_max_suppression,
+ scale_boxes,
+ xywh2xyxy,
+ xyxy2xywh,
+)
+from yolov5.utils.metrics import ap_per_class, box_iou
+
+from sparsebit.quantization import QuantModel, parse_qconfig
+
+
+def set_seed(seed):
+ os.environ["PYTHONHASHSEED"] = str(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+ torch.backends.cudnn.enabled = False
+
+
+def save_one_txt(predn, save_conf, shape, file):
+ # Save one txt result
+ gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
+ for *xyxy, conf, cls in predn.tolist():
+ xywh = (
+ (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
+ ) # normalized xywh
+ line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
+ with open(file, "a") as f:
+ f.write(("%g " * len(line)).rstrip() % line + "\n")
+
+
+def save_one_json(predn, jdict, path, class_map):
+ # Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
+ image_id = int(path.stem) if path.stem.isnumeric() else path.stem
+ box = xyxy2xywh(predn[:, :4]) # xywh
+ box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
+ for p, b in zip(predn.tolist(), box.tolist()):
+ jdict.append(
+ {
+ "image_id": image_id,
+ "category_id": class_map[int(p[5])],
+ "bbox": [round(x, 3) for x in b],
+ "score": round(p[4], 5),
+ }
+ )
+
+
+def process_batch(detections, labels, iouv):
+ """
+ Return correct prediction matrix
+ Arguments:
+ detections (array[N, 6]), x1, y1, x2, y2, conf, class
+ labels (array[M, 5]), class, x1, y1, x2, y2
+ Returns:
+ correct (array[N, 10]), for 10 IoU levels
+ """
+ correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
+ iou = box_iou(labels[:, 1:], detections[:, :4])
+ correct_class = labels[:, 0:1] == detections[:, 5]
+ for i in range(len(iouv)):
+ x = torch.where(
+ (iou >= iouv[i]) & correct_class
+ ) # IoU > threshold and classes match
+ if x[0].shape[0]:
+ matches = (
+ torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1)
+ .cpu()
+ .numpy()
+ ) # [label, detect, iou]
+ if x[0].shape[0] > 1:
+ matches = matches[matches[:, 2].argsort()[::-1]]
+ matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
+ # matches = matches[matches[:, 2].argsort()[::-1]]
+ matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
+ correct[matches[:, 1].astype(int), i] = True
+ return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
+
+
+@torch.no_grad()
+def main(args):
+ set_seed(args.seed)
+ if args.model_name == "yolov5n":
+ model = yolov5n(checkpoint_path=args.checkpoint_path)
+ elif args.model_name == "yolov5s":
+ model = yolov5s(checkpoint_path=args.checkpoint_path)
+ else:
+ raise NotImplementedError
+
+ qconfig = parse_qconfig(args.qconfig_path)
+ qmodel = QuantModel(model.model, config=qconfig)
+
+ setattr(model, "model", qmodel)
+ imgsz = 640
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ # Configure
+ model.eval()
+ model.cuda()
+
+ cuda = device.type != "cpu"
+ nc = 80 # number of classes
+ iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for [email protected]:0.95
+ niou = iouv.numel()
+ task = "val" # path to val images
+ dataloader = create_dataloader(
+ os.path.join(args.data_path, "val2017.txt"),
+ imgsz,
+ 1,
+ 32,
+ False,
+ pad=0.5,
+ rect=True,
+ workers=args.workers,
+ prefix=colorstr(f"{task}: "),
+ )[0]
+
+ # Calibration
+ calib_loader = create_dataloader(
+ os.path.join(args.data_path, "calib2017.txt"),
+ imgsz,
+ 1,
+ 32,
+ False,
+ pad=0.5,
+ rect=True,
+ workers=args.workers,
+ prefix=colorstr(f"{task}: "),
+ )[0]
+
+ qmodel.prepare_calibration()
+ for batch_meta in calib_loader:
+ data = batch_meta[0] / 255
+ with torch.no_grad():
+ _ = qmodel(data.to(device, non_blocking=True))
+ qmodel.calc_qparams()
+ qmodel.set_quant(w_quant=True, a_quant=True)
+
+ seen = 0
+ 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))
+ s = ("%22s" + "%11s" * 6) % (
+ "Class",
+ "Images",
+ "Instances",
+ "P",
+ "R",
+ "mAP50",
+ "mAP50-95",
+ )
+ p, r, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
+ dt = Profile(), Profile(), Profile() # profiling times
+ stats, ap = [], []
+ pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT) # progress bar
+ for im, targets, paths, shapes in pbar:
+ with dt[0]:
+ if cuda:
+ im = im.to(device, non_blocking=True)
+ targets = targets.to(device)
+ im = im.float() # uint8 to fp16/32
+ im /= 255 # 0 - 255 to 0.0 - 1.0
+ _, _, height, width = im.shape # batch size, channels, height, width
+
+ # Inference
+ with dt[1]:
+ preds = model(im)
+
+ # NMS
+ targets[:, 2:] *= torch.tensor(
+ (width, height, width, height), device=device
+ ) # to pixels
+ with dt[2]:
+ preds = non_max_suppression(
+ preds,
+ args.conf_thres,
+ args.iou_thres,
+ labels=[],
+ multi_label=True,
+ agnostic=False,
+ max_det=300,
+ )
+
+ # Metrics
+ for si, pred in enumerate(preds):
+ labels = targets[targets[:, 0] == si, 1:]
+ npr = pred.shape[0] # number of labels, predictions
+ shape = shapes[si][0]
+ correct = torch.zeros(npr, niou, dtype=torch.bool, device=device) # init
+ seen += 1
+
+ if npr == 0:
+ stats.append(
+ (correct, *torch.zeros((2, 0), device=device), labels[:, 0])
+ )
+ continue
+
+ # Predictions
+ predn = pred.clone()
+ scale_boxes(
+ im[si].shape[1:], predn[:, :4], shape, shapes[si][1]
+ ) # native-space pred
+
+ # Evaluate
+ tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
+ scale_boxes(
+ im[si].shape[1:], tbox, shape, shapes[si][1]
+ ) # native-space labels
+ labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
+ correct = process_batch(predn, labelsn, iouv)
+
+ stats.append(
+ (correct, pred[:, 4], pred[:, 5], labels[:, 0])
+ ) # (correct, conf, pcls, tcls)
+
+ # Compute metrics
+ stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)] # to numpy
+ if len(stats) and stats[0].any():
+ tp, fp, p, r, f1, ap, ap_class = ap_per_class(
+ *stats, plot=False, save_dir="", names=names
+ )
+ 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
+
+ # Print results
+ pf = "%22s" + "%11i" * 2 + "%11.3g" * 4 # print format
+ LOGGER.info(pf % ("all", seen, nt.sum(), mp, mr, map50, map))
+ if nt.sum() == 0:
+ LOGGER.warning(
+ f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels"
+ )
+
+
+def parse_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--model_name", type=str, default="yolov5n", choices=["yolov5n", "yolov5s"]
+ )
+ parser.add_argument("--qconfig_path", type=str, default="./qconfig.yaml")
+ parser.add_argument("--data_path", type=str, required=True)
+ parser.add_argument("--checkpoint_path", type=str, required=True)
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--workers", type=int, default=8, help="dataloader workers")
+ 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"
+ )
+
+ args = parser.parse_args()
+ return args
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ main(args)