Refactor the code

.github/ISSUE_TEMPLATE/sweep-template.yml

@@ -0,0 +1,15 @@
+name: Sweep Issue
+title: 'Sweep: '
+description: For small bugs, features, refactors, and tests to be handled by Sweep, an AI-powered junior developer.
+labels: sweep
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Details
+      description: Tell Sweep where and what to edit and provide enough context for a new developer to the codebase
+      placeholder: |
+        Unit Tests: Write unit tests for <FILE>. Test each function in the file. Make sure to test edge cases.
+        Bugs: The bug might be in <FILE>. Here are the logs: ...
+        Features: the new endpoint should use the ... class from <FILE> because it contains ... logic.
+        Refactors: We are migrating this function to ... version because ...

README.md

@@ -2,6 +2,8 @@ Quantization Aware Training Implementation of YOLOv8 without [DFL](https://ieeex
 
 ### Installation
 
+Execute the command:
+
 ```
 conda create -n YOLO python=3.8
 conda activate YOLO
@@ -11,6 +13,12 @@ pip install PyYAML
 pip install tqdm
 ```
 
+or 
+
+```
+pip3 install -r requirements.txt
+```
+
 ### Train
 
 * Configure your dataset path in `main.py` for training

main.py

@@ -1,287 +1,25 @@
-import copy
-import csv
 import os
 import warnings
 from argparse import ArgumentParser
 
-import numpy
 import torch
-import tqdm
-import yaml
-from torch.utils import data
+from torch.utils.data import DataLoader
+from torchvision import transforms
 
 from nets import nn
 from utils import util
-from utils.dataset import Dataset
+from utils.dataset import CustomDataset
+from utils.trainer import Trainer
+from utils.tester import Tester
 
 warnings.filterwarnings("ignore")
 
 
-def learning_rate(args, params):
-    def fn(x):
-        return (1 - x / args.epochs) * (1.0 - params['lrf']) + params['lrf']
-
-    return fn
-
-
-def train(args, params):
-    util.setup_seed()
-    util.setup_multi_processes()
-
-    # Model
-    model = nn.yolo_v8_n(len(params['names']))
-    state = torch.load('./weights/v8_n.pth')['model']
-    model.load_state_dict(state.float().state_dict())
-    model.eval()
-
-    for m in model.modules():
-        if type(m) is nn.Conv and hasattr(m, 'norm'):
-            torch.ao.quantization.fuse_modules(m, [["conv", "norm"]], True)
-    model.train()
-
-    model = nn.QAT(model)
-    model.qconfig = torch.quantization.get_default_qconfig("qnnpack")
-    torch.quantization.prepare_qat(model, inplace=True)
-    model.cuda()
-
-    # Optimizer
-    accumulate = max(round(64 / (args.batch_size * args.world_size)), 1)
-    params['weight_decay'] *= args.batch_size * args.world_size * accumulate / 64
-
-    optimizer = torch.optim.SGD(util.weight_decay(model, params['weight_decay']),
-                                params['lr0'], params['momentum'], nesterov=True)
-
-    # Scheduler
-    lr = learning_rate(args, params)
-    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr, last_epoch=-1)
-
-    filenames = []
-    with open('../Dataset/COCO/train2017.txt') as reader:
-        for filename in reader.readlines():
-            filename = filename.rstrip().split('/')[-1]
-            filenames.append('../Dataset/COCO/images/train2017/' + filename)
-
-    sampler = None
-    dataset = Dataset(filenames, args.input_size, params, True)
-
-    if args.distributed:
-        sampler = data.distributed.DistributedSampler(dataset)
-
-    loader = data.DataLoader(dataset, args.batch_size, sampler is None, sampler,
-                             num_workers=8, pin_memory=True, collate_fn=Dataset.collate_fn)
-
-    if args.distributed:
-        # DDP mode
-        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
-        model = torch.nn.parallel.DistributedDataParallel(module=model,
-                                                          device_ids=[args.local_rank],
-                                                          output_device=args.local_rank)
-
-    best = 0
-    num_steps = len(loader)
-    criterion = util.ComputeLoss(model, params)
-    num_warmup = max(round(params['warmup_epochs'] * num_steps), 100)
-    with open('weights/step.csv', 'w') as f:
-        if args.local_rank == 0:
-            writer = csv.DictWriter(f, fieldnames=['epoch',
-                                                   'box', 'cls',
-                                                   'Recall', 'Precision', 'mAP@50', 'mAP'])
-            writer.writeheader()
-        for epoch in range(args.epochs):
-            model.train()
-            if args.distributed:
-                sampler.set_epoch(epoch)
-            if args.epochs - epoch == 10:
-                loader.dataset.mosaic = False
-
-            p_bar = enumerate(loader)
-
-            if args.local_rank == 0:
-                print(('\n' + '%10s' * 4) % ('epoch', 'memory', 'box', 'cls'))
-            if args.local_rank == 0:
-                p_bar = tqdm.tqdm(p_bar, total=num_steps)  # progress bar
-
-            optimizer.zero_grad()
-            avg_box_loss = util.AverageMeter()
-            avg_cls_loss = util.AverageMeter()
-            for i, (samples, targets) in p_bar:
-                samples = samples.cuda()
-                samples = samples.float()
-                samples = samples / 255.0
-
-                x = i + num_steps * epoch
-
-                # Warmup
-                if x <= num_warmup:
-                    xp = [0, num_warmup]
-                    fp = [1, 64 / (args.batch_size * args.world_size)]
-                    accumulate = max(1, numpy.interp(x, xp, fp).round())
-                    for j, y in enumerate(optimizer.param_groups):
-                        if j == 0:
-                            fp = [params['warmup_bias_lr'], y['initial_lr'] * lr(epoch)]
-                        else:
-                            fp = [0.0, y['initial_lr'] * lr(epoch)]
-                        y['lr'] = numpy.interp(x, xp, fp)
-                        if 'momentum' in y:
-                            fp = [params['warmup_momentum'], params['momentum']]
-                            y['momentum'] = numpy.interp(x, xp, fp)
-
-                # Forward
-                outputs = model(samples)
-                loss_box, loss_cls = criterion(outputs, targets)
-
-                avg_box_loss.update(loss_box.item(), samples.size(0))
-                avg_cls_loss.update(loss_cls.item(), samples.size(0))
-
-                loss_box *= args.batch_size  # loss scaled by batch_size
-                loss_cls *= args.batch_size  # loss scaled by batch_size
-                loss_box *= args.world_size  # gradient averaged between devices in DDP mode
-                loss_cls *= args.world_size  # gradient averaged between devices in DDP mode
-
-                # Backward
-                (loss_box + loss_cls).backward()
-
-                # Optimize
-                if x % accumulate == 0:
-                    util.clip_gradients(model)  # clip gradients
-                    optimizer.step()
-                    optimizer.zero_grad()
-
-                # Log
-                if args.local_rank == 0:
-                    memory = f'{torch.cuda.memory_reserved() / 1E9:.4g}G'  # (GB)
-                    s = ('%10s' * 2 + '%10.3g' * 2) % (f'{epoch + 1}/{args.epochs}', memory,
-                                                       avg_box_loss.avg, avg_cls_loss.avg)
-                    p_bar.set_description(s)
-
-            # Scheduler
-            scheduler.step()
-
-            if args.local_rank == 0:
-                # Convert model
-                save = copy.deepcopy(model.module if args.distributed else model)
-                save.eval()
-                save.to(torch.device('cpu'))
-                torch.ao.quantization.convert(save, inplace=True)
-                # mAP
-                last = test(args, params, save)
-
-                writer.writerow({'epoch': str(epoch + 1).zfill(3),
-                                 'box': str(f'{avg_box_loss.avg:.3f}'),
-                                 'cls': str(f'{avg_cls_loss.avg:.3f}'),
-                                 'mAP': str(f'{last[0]:.3f}'),
-                                 'mAP@50': str(f'{last[1]:.3f}'),
-                                 'Recall': str(f'{last[2]:.3f}'),
-                                 'Precision': str(f'{last[2]:.3f}')})
-                f.flush()
-
-                # Update best mAP
-                if last[0] > best:
-                    best = last[0]
-
-                # Save last, best and delete
-                save = torch.jit.script(save.cpu())
-                torch.jit.save(save, './weights/last.ts')
-                if best == last[0]:
-                    torch.jit.save(save, './weights/best.ts')
-                del save
-
-    torch.cuda.empty_cache()
-
-
-@torch.no_grad()
-def test(args, params, model=None):
-    filenames = []
-    with open('../Dataset/COCO/val2017.txt') as reader:
-        for filename in reader.readlines():
-            filename = filename.rstrip().split('/')[-1]
-            filenames.append('../Dataset/COCO/images/val2017/' + filename)
-
-    dataset = Dataset(filenames, args.input_size, params, False)
-    loader = data.DataLoader(dataset, args.batch_size // 2, False, num_workers=8,
-                             pin_memory=True, collate_fn=Dataset.collate_fn)
-    if model is None:
-        model = torch.jit.load(f='./weights/best.ts')
-
-    device = torch.device('cpu')
-    model.to(device)
-    model.eval()
-
-    # Configure
-    iou_v = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for mAP@0.5:0.95
-    n_iou = iou_v.numel()
-
-    m_pre = 0.
-    m_rec = 0.
-    map50 = 0.
-    mean_ap = 0.
-    metrics = []
-    p_bar = tqdm.tqdm(loader, desc=('%10s' * 4) % ('precision', 'recall', 'mAP50', 'mAP'))
-    for samples, targets in p_bar:
-        samples = samples.to(device)
-        samples = samples.float()  # uint8 to fp16/32
-        samples = samples / 255.0  # 0 - 255 to 0.0 - 1.0
-        _, _, h, w = samples.shape  # batch size, channels, height, width
-        scale = torch.tensor((w, h, w, h), device=device)
-        # Inference
-        outputs = model(samples)
-        # NMS
-        outputs = util.non_max_suppression(outputs, 0.001, 0.7, model.nc)
-        # Metrics
-        for i, output in enumerate(outputs):
-            idx = targets['idx'] == i
-            cls = targets['cls'][idx]
-            box = targets['box'][idx]
-
-            cls = cls.to(device)
-            box = box.to(device)
-
-            metric = torch.zeros(output.shape[0], n_iou, dtype=torch.bool, device=device)
-
-            if output.shape[0] == 0:
-                if cls.shape[0]:
-                    metrics.append((metric, *torch.zeros((2, 0), device=device), cls.squeeze(-1)))
-                continue
-            # Evaluate
-            if cls.shape[0]:
-                target = torch.cat((cls, util.wh2xy(box) * scale), 1)
-                metric = util.compute_metric(output[:, :6], target, iou_v)
-            # Append
-            metrics.append((metric, output[:, 4], output[:, 5], cls.squeeze(-1)))
-
-    # Compute metrics
-    metrics = [torch.cat(x, 0).cpu().numpy() for x in zip(*metrics)]  # to numpy
-    if len(metrics) and metrics[0].any():
-        tp, fp, m_pre, m_rec, map50, mean_ap = util.compute_ap(*metrics)
-    # Print results
-    print('%10.3g' * 4 % (m_pre, m_rec, map50, mean_ap))
-    # Return results
-    model.float()  # for training
-    return mean_ap, map50, m_rec, m_pre
-
-
-def profile(args, params):
-    from thop import profile, clever_format
-    model = nn.yolo_v8_n(len(params['names']))
-    shape = (1, 3, args.input_size, args.input_size)
-
-    model.eval()
-    torch.set_num_threads(1)
-    torch.set_num_interop_threads(1)
-
-    macs, params = profile(model, inputs=(torch.zeros(shape),), verbose=False)
-    macs, params = clever_format([macs, params], "%.3f")
-
-    if args.local_rank == 0:
-        print(f'MACs: {macs}')
-        print(f'Parameters: {params}')
-
-
 def main():
     parser = ArgumentParser()
     parser.add_argument('--input-size', default=640, type=int)
     parser.add_argument('--batch-size', default=32, type=int)
-    parser.add_argument('--local_rank', default=0, type=int)
+    parser.add_argument('--local-rank', default=0, type=int)
     parser.add_argument('--epochs', default=20, type=int)
     parser.add_argument('--train', action='store_true')
     parser.add_argument('--test', action='store_true')
@@ -300,13 +38,17 @@ def main():
         if not os.path.exists('weights'):
             os.makedirs('weights')
 
-    with open('utils/args.yaml', errors='ignore') as f:
-        params = yaml.safe_load(f)
-    profile(args, params)
+    config_path = 'utils/args.yaml'
+    params = util.load_config(config_path)
+
     if args.train:
-        train(args, params)
+        trainer = Trainer(args, params)
+        best_mean_ap = trainer.train()
+        print(f'Best mAP: {best_mean_ap:.3f}')
+
     if args.test:
-        test(args, params)
+        tester = Tester(args, params)
+        tester.test()
 
 
 if __name__ == "__main__":

requirements.txt

@@ -0,0 +1,7 @@
+albumentations==1.3.1
+numpy==1.24.3
+Pillow==10.2.0
+PyYAML==6.0.1
+torch==2.1.2
+torchvision==0.16.2
+tqdm==4.65.0

sweep.yaml

@@ -0,0 +1,27 @@
+# Sweep AI turns bugs & feature requests into code changes (https://sweep.dev)
+# For details on our config file, check out our docs at https://docs.sweep.dev/usage/config
+
+# This setting contains a list of rules that Sweep will check for. If any of these rules are broken in a new commit, Sweep will create an pull request to fix the broken rule.
+rules:
+  - "All new business logic should have corresponding unit tests."
+  - "Refactor large functions to be more modular."
+  - "Add docstrings to all functions and file headers."
+
+# This is the branch that Sweep will develop from and make pull requests to. Most people use 'main' or 'master' but some users also use 'dev' or 'staging'.
+branch: 'main'
+
+# By default Sweep will read the logs and outputs from your existing Github Actions. To disable this, set this to false.
+gha_enabled: True
+
+# This is the description of your project. It will be used by sweep when creating PRs. You can tell Sweep what's unique about your project, what frameworks you use, or anything else you want.
+#
+# Example:
+#
+# description: sweepai/sweep is a python project. The main api endpoints are in sweepai/api.py. Write code that adheres to PEP8.
+description: ''
+
+# This sets whether to create pull requests as drafts. If this is set to True, then all pull requests will be created as drafts and GitHub Actions will not be triggered.
+draft: False
+
+# This is a list of directories that Sweep will not be able to edit.
+blocked_dirs: []