Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor the code #6

Open
wants to merge 28 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
728f007
Refactor the codes
yihong1120 Jan 11, 2024
43a97d3
Refactor and optimised the code
yihong1120 Jan 11, 2024
bceb7fa
Refactor and optimised the code
yihong1120 Jan 11, 2024
93b95b2
Add required packages
yihong1120 Jan 11, 2024
43fd81d
Add new package installation pattern
yihong1120 Jan 11, 2024
1879fc3
Create LICENSE
yihong1120 Jan 11, 2024
121e208
Create sweep.yaml
sweep-ai[bot] Jan 12, 2024
b04206b
Create sweep template
sweep-ai[bot] Jan 12, 2024
4e51c28
Merge pull request #3 from yihong1120/sweep/add-sweep-config
yihong1120 Jan 12, 2024
1bb7e56
feat: Add unit tests for main.py
sweep-ai[bot] Jan 12, 2024
c99960e
feat: Add unit tests for dataset.py
sweep-ai[bot] Jan 12, 2024
918eb26
feat: Add unit tests for tester.py
sweep-ai[bot] Jan 12, 2024
fcf369c
feat: Add unit tests for trainer.py
sweep-ai[bot] Jan 12, 2024
cc615d2
feat: Add unit tests for util.py
sweep-ai[bot] Jan 12, 2024
e225d98
feat: Updated tests/test_main.py
sweep-ai[bot] Jan 12, 2024
d98f3c9
feat: Updated tests/test_dataset.py
sweep-ai[bot] Jan 12, 2024
c32e603
feat: Updated tests/test_tester.py
sweep-ai[bot] Jan 12, 2024
893d990
feat: Updated tests/test_trainer.py
sweep-ai[bot] Jan 12, 2024
948fa12
feat: Updated tests/test_util.py
sweep-ai[bot] Jan 12, 2024
32d0bd6
Merge pull request #4 from yihong1120/sweep/create_files_according_to…
yihong1120 Jan 12, 2024
6e58a50
Complete code
yihong1120 Jan 13, 2024
7165a0a
Complete code
yihong1120 Jan 13, 2024
2cdc8e6
Add type hints
yihong1120 Jan 13, 2024
a19a1e6
Add type hints
yihong1120 Jan 13, 2024
9e15d63
Complete code
yihong1120 Jan 13, 2024
5d364b1
Add docstring
yihong1120 Jan 13, 2024
ef2b3ee
Add docstring
yihong1120 Jan 13, 2024
752c3b0
Complete code
yihong1120 Jan 13, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions .github/ISSUE_TEMPLATE/sweep-template.yml
Original file line number Diff line number Diff line change
@@ -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 ...
661 changes: 661 additions & 0 deletions LICENSE

Large diffs are not rendered by default.

8 changes: 8 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
288 changes: 15 additions & 273 deletions main.py
Original file line number Diff line number Diff line change
@@ -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 [email protected]: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')
Expand All @@ -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__":
Expand Down
7 changes: 7 additions & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -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
27 changes: 27 additions & 0 deletions sweep.yaml
Original file line number Diff line number Diff line change
@@ -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: []
Loading