From a217cd390a5175e14367221058cce4a33a3b3356 Mon Sep 17 00:00:00 2001 From: BBuf <1182563586@qq.com> Date: Wed, 14 Dec 2022 14:19:38 +0000 Subject: [PATCH] add build_targets_optim --- train.py | 6 ++++-- utils/loss.py | 12 ++++++++---- utils/metrics.py | 14 +++++++------- 3 files changed, 19 insertions(+), 13 deletions(-) diff --git a/train.py b/train.py index 56c87b93..00d50ba8 100644 --- a/train.py +++ b/train.py @@ -84,7 +84,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary - (save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, bbox_iou_optim, multi_tensor_optimizer) = ( + (save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, bbox_iou_optim, build_targets_optim, multi_tensor_optimizer) = ( Path(opt.save_dir), opt.epochs, opt.batch_size, @@ -99,6 +99,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio opt.workers, opt.freeze, opt.bbox_iou_optim, + opt.build_targets_optim, opt.multi_tensor_optimizer, ) @@ -285,7 +286,7 @@ def f(x): # scaler = flow.cuda.amp.GradScaler(enabled=amp) stopper, _ = EarlyStopping(patience=opt.patience), False - compute_loss = ComputeLoss(model, bbox_iou_optim=bbox_iou_optim) # init loss class + compute_loss = ComputeLoss(model, bbox_iou_optim=bbox_iou_optim, build_targets_optim=build_targets_optim) # init loss class callbacks.run("on_train_start") LOGGER.info( f"Image sizes {imgsz} train, {imgsz} val\n" @@ -517,6 +518,7 @@ def parse_opt(known=False): 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("--bbox_iou_optim", action="store_true", help="optim bbox_iou function in compute_loss") + parser.add_argument("--build_targets_optim", action="store_true", help="optim build_targets function in compute_loss") parser.add_argument("--multi_tensor_optimizer", action="store_true", help="apply multi_tensor implement in optimizer") parser.add_argument( "--patience", diff --git a/utils/loss.py b/utils/loss.py index 1a317af9..5712f16c 100644 --- a/utils/loss.py +++ b/utils/loss.py @@ -93,7 +93,7 @@ class ComputeLoss: sort_obj_iou = False # Compute losses - def __init__(self, model, autobalance=False, bbox_iou_optim=False): + def __init__(self, model, autobalance=False, bbox_iou_optim=False, build_targets_optim=False): device = next(model.parameters()).device # get model device h = model.hyp # hyperparameters @@ -125,6 +125,7 @@ def __init__(self, model, autobalance=False, bbox_iou_optim=False): self.anchors = m.anchors self.device = device self.bbox_iou_optim = bbox_iou_optim + self.build_targets_optim = build_targets_optim def __call__(self, p, targets): # predictions, targets lcls = flow.zeros(1, device=self.device) # class loss @@ -235,9 +236,12 @@ def build_targets(self, p, targets): # Offsets gxy = t[:, 2:4] # grid xy gxi = gain[[2, 3]] - gxy # inverse - j, k = ((gxy % 1 < g) & (gxy > 1)).T - l, m = ((gxi % 1 < g) & (gxi > 1)).T - j = flow.stack((flow.ones_like(j), j, k, l, m)) + if self.build_targets_optim: + j = flow._C.fused_yolov5_get_target_offsets(gxy, gxi, g) + else: + j, k = ((gxy % 1 < g) & (gxy > 1)).T + l, m = ((gxi % 1 < g) & (gxi > 1)).T + j = flow.stack((flow.ones_like(j), j, k, l, m)) t = t.repeat((5, 1, 1))[j] offsets = (flow.zeros_like(gxy)[None] + off[:, None])[j] else: diff --git a/utils/metrics.py b/utils/metrics.py index f93d1ffa..d1815e52 100644 --- a/utils/metrics.py +++ b/utils/metrics.py @@ -228,7 +228,7 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, bbox_iou if xywh: # transform from xywh to xyxy (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, 1), box2.chunk(4, 1) if bbox_iou_optim: - b1_x1, b1_x2, b1_y1, b1_y2, b2_x1, b2_x2, b2_y1, b2_y2 = flow._C.fused_get_boundding_boxes_coord(x1, y1, w1, h1, x2, y2, w2, h2) + b1_x1, b1_x2, b1_y1, b1_y2, b2_x1, b2_x2, b2_y1, b2_y2 = flow._C.fused_yolov5_get_boundding_boxes_coord(x1, y1, w1, h1, x2, y2, w2, h2) else: w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2 b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_ @@ -241,12 +241,12 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, bbox_iou # Intersection area if bbox_iou_optim: - inter = flow._C.fused_get_intersection_area(b1_x1, b1_x2, b2_x1, b2_x2, b1_y1, b1_y2, b2_y1, b2_y2) + inter = flow._C.fused_yolov5_get_intersection_area(b1_x1, b1_x2, b2_x1, b2_x2, b1_y1, b1_y2, b2_y1, b2_y2) else: inter = (flow.min(b1_x2, b2_x2) - flow.max(b1_x1, b2_x1)).clamp(0) * (flow.min(b1_y2, b2_y2) - flow.max(b1_y1, b2_y1)).clamp(0) if bbox_iou_optim and CIoU: - iou = flow._C.fused_get_iou(w1, h1, w2, h2, inter, eps) + iou = flow._C.fused_yolov5_get_iou(w1, h1, w2, h2, inter, eps) else: # Union Area union = w1 * h1 + w2 * h2 - inter + eps @@ -260,18 +260,18 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, bbox_iou ch = flow.max(b1_y2, b2_y2) - flow.min(b1_y1, b2_y1) # convex height if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 if bbox_iou_optim: - c2 = flow._C.fused_get_convex_diagonal_squared(b1_x1, b1_x2, b2_x1, b2_x2, b1_y1, b1_y2, b2_y1, b2_y2, eps) + c2 = flow._C.fused_yolov5_get_convex_diagonal_squared(b1_x1, b1_x2, b2_x1, b2_x2, b1_y1, b1_y2, b2_y1, b2_y2, eps) else: c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared if bbox_iou_optim: - rho2 = flow._C.fused_get_center_dist(b1_x1, b1_x2, b2_x1, b2_x2, b1_y1, b1_y2, b2_y1, b2_y2) + rho2 = flow._C.fused_yolov5_get_center_dist(b1_x1, b1_x2, b2_x1, b2_x2, b1_y1, b1_y2, b2_y1, b2_y2) else: rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center dist ** 2 if CIoU: if bbox_iou_optim: - v = flow._C.fused_get_ciou_diagonal_angle(w1, h1, w2, h2, eps) - return flow._C.fused_get_ciou_result(v, iou, rho2, c2, eps)[0] + v = flow._C.fused_yolov5_get_ciou_diagonal_angle(w1, h1, w2, h2, eps) + return flow._C.fused_yolov5_get_ciou_result(v, iou, rho2, c2, eps)[0] else: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 v = (4 / math.pi ** 2) * flow.pow(flow.atan(w2 / (h2 + eps)) - flow.atan(w1 / (h1 + eps)), 2)