New YOLOv5 Classification Models (#8956)
* Update * Logger step fix: Increment step with epochs (#8654) * enhance * revert * allow training from scratch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update --img argument from train.py single line * fix image size from 640 to 128 * suport custom dataloader and augmentation * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * format * Update dataloaders.py * Single line return, single line comment, remove unused argument * address PR comments * fix spelling * don't augment eval set * use fstring * update augmentations.py * new maning convention for transforms * reverse if statement, inline ops * reverse if statement, inline ops * updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update dataloaders * Remove additional if statement * Remove is_train as redundant * Cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup2 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update classifier.py * Update augmentations.py * fix: imshow clip warning * update * Revert ToTensorV2 removal * Update classifier.py * Update normalize values, revert uint8 * normalize image using cv2 * remove dedundant comment * Update classifier.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * replace print with logger * commit steps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Allow logging models from GenericLogger (#8676) * enhance * revert * allow training from scratch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update --img argument from train.py single line * fix image size from 640 to 128 * suport custom dataloader and augmentation * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * format * Update dataloaders.py * Single line return, single line comment, remove unused argument * address PR comments * fix spelling * don't augment eval set * use fstring * update augmentations.py * new maning convention for transforms * reverse if statement, inline ops * reverse if statement, inline ops * updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update dataloaders * Remove additional if statement * Remove is_train as redundant * Cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup2 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update classifier.py * Update augmentations.py * fix: imshow clip warning * update * Revert ToTensorV2 removal * Update classifier.py * Update normalize values, revert uint8 * normalize image using cv2 * remove dedundant comment * Update classifier.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * replace print with logger * commit steps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * support final model logging * update * update * update * update * remove curses * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update classifier.py * Update __init__.py Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Update * Update * Update * Update * Update dataset download * Update dataset download * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Pass imgsz to classify_transforms() * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Cos scheduler * Cos scheduler * Remove unused args * Update * Add seed * Add seed * Update * Update * Add run(), main() * Merge master * Merge master * Update * Update * Update * Update * Update * Update * Update * Create YOLOv5 BaseModel class (#8829) * Create BaseModel * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Hub load device fix * Update Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Add experiment * Merge master * Attach names * weight decay = 1e-4 * weight decay = 5e-5 * update smart_optimizer console printout * fashion-mnist fix * Merge master * Update Table * Update Table * Remove destroy process group * add kwargs to forward() * fuse fix for resnet50 * nc, names fix for resnet50 * nc, names fix for resnet50 * ONNX CPU inference fix * revert * cuda * if augment or visualize * if augment or visualize * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * New smart_inference_mode() * Update README * Refactor into /classify dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reset defaults * reset defaults * fix gpu predict * warmup * ema half fix * spacing * remove data * remove cache * remove denormalize * save run settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * verbose false on initial plots * new save_yaml() function * Update ci-testing.yml * Path(data) CI fix * Separate classification CI * fix val * fix val * fix val * smartCrossEntropyLoss * skip validation on hub load * autodownload with working dir root * str(data) * Dataset usage example * im_show normalize * im_show normalize * add imagenet simple names to multibackend * Add validation speeds * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 24-space names * Update bash scripts * Update permissions * Add bash script arguments * remove verbose * TRT data fix * names generator fix * optimize if names * update usage * Add local loading * Verbose=False * update names printing * Add Usage examples * Add Usage examples * Add Usage examples * Add Usage examples * named_children * reshape_classifier_outputs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update * update * fix CI * fix incorrect class substitution * fix incorrect class substitution * remove denormalize * ravel fix * cleanup * update opt file printing * update opt file printing * update defaults * add opt to checkpoint * Add warning * Add comment * plot half bug fix * Use NotImplementedError * fix export shape report * Fix TRT load * cleanup CI * profile comment * CI fix * Add cls models * avoid inplace error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix usage examples * Update README * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README * Update README Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
22 files changed
+1436
-161
No files matched your search
@@ -0,0 +1,109 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Run classification inference on images
|
||||
|
||||
Usage:
|
||||
$ python classify/predict.py --weights yolov5s-cls.pt --source im.jpg
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import torch.nn.functional as F
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[1] # YOLOv5 root directory
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||
|
||||
from classify.train import imshow_cls
|
||||
from models.common import DetectMultiBackend
|
||||
from utils.augmentations import classify_transforms
|
||||
from utils.general import LOGGER, check_requirements, colorstr, increment_path, print_args
|
||||
from utils.torch_utils import select_device, smart_inference_mode, time_sync
|
||||
|
||||
|
||||
@smart_inference_mode()
|
||||
def run(
|
||||
weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s)
|
||||
source=ROOT / 'data/images/bus.jpg', # file/dir/URL/glob, 0 for webcam
|
||||
imgsz=224, # inference size
|
||||
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||
half=False, # use FP16 half-precision inference
|
||||
dnn=False, # use OpenCV DNN for ONNX inference
|
||||
show=True,
|
||||
project=ROOT / 'runs/predict-cls', # save to project/name
|
||||
name='exp', # save to project/name
|
||||
exist_ok=False, # existing project/name ok, do not increment
|
||||
):
|
||||
file = str(source)
|
||||
seen, dt = 1, [0.0, 0.0, 0.0]
|
||||
device = select_device(device)
|
||||
|
||||
# Directories
|
||||
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
||||
save_dir.mkdir(parents=True, exist_ok=True) # make dir
|
||||
|
||||
# Transforms
|
||||
transforms = classify_transforms(imgsz)
|
||||
|
||||
# Load model
|
||||
model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
|
||||
model.warmup(imgsz=(1, 3, imgsz, imgsz)) # warmup
|
||||
|
||||
# Image
|
||||
t1 = time_sync()
|
||||
im = cv2.cvtColor(cv2.imread(file), cv2.COLOR_BGR2RGB)
|
||||
im = transforms(im).unsqueeze(0).to(device)
|
||||
im = im.half() if model.fp16 else im.float()
|
||||
t2 = time_sync()
|
||||
dt[0] += t2 - t1
|
||||
|
||||
# Inference
|
||||
results = model(im)
|
||||
t3 = time_sync()
|
||||
dt[1] += t3 - t2
|
||||
|
||||
p = F.softmax(results, dim=1) # probabilities
|
||||
i = p.argsort(1, descending=True)[:, :5].squeeze() # top 5 indices
|
||||
dt[2] += time_sync() - t3
|
||||
LOGGER.info(f"image 1/1 {file}: {imgsz}x{imgsz} {', '.join(f'{model.names[j]} {p[0, j]:.2f}' for j in i)}")
|
||||
|
||||
# Print results
|
||||
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
|
||||
shape = (1, 3, imgsz, imgsz)
|
||||
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}' % t)
|
||||
if show:
|
||||
imshow_cls(im, f=save_dir / Path(file).name, verbose=True)
|
||||
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
|
||||
return p
|
||||
|
||||
|
||||
def parse_opt():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)')
|
||||
parser.add_argument('--source', type=str, default=ROOT / 'data/images/bus.jpg', help='file')
|
||||
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='train, val image size (pixels)')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
||||
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
|
||||
parser.add_argument('--project', default=ROOT / 'runs/predict-cls', 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')
|
||||
opt = parser.parse_args()
|
||||
print_args(vars(opt))
|
||||
return opt
|
||||
|
||||
|
||||
def main(opt):
|
||||
check_requirements(exclude=('tensorboard', 'thop'))
|
||||
run(**vars(opt))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
opt = parse_opt()
|
||||
main(opt)
|
||||
@@ -0,0 +1,325 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Train a YOLOv5 classifier model on a classification dataset
|
||||
Datasets: --data mnist, fashion-mnist, cifar10, cifar100, imagenette, imagewoof, imagenet, or 'path/to/custom/dataset'
|
||||
|
||||
Usage:
|
||||
$ python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 128
|
||||
$ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from copy import deepcopy
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.hub as hub
|
||||
import torch.optim.lr_scheduler as lr_scheduler
|
||||
import torchvision
|
||||
from torch.cuda import amp
|
||||
from tqdm import tqdm
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[1] # YOLOv5 root directory
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||
|
||||
from classify import val as validate
|
||||
from models.experimental import attempt_load
|
||||
from models.yolo import ClassificationModel, DetectionModel
|
||||
from utils.dataloaders import create_classification_dataloader
|
||||
from utils.general import (DATASETS_DIR, LOGGER, WorkingDirectory, check_git_status, check_requirements, colorstr,
|
||||
download, increment_path, init_seeds, print_args, yaml_save)
|
||||
from utils.loggers import GenericLogger
|
||||
from utils.plots import imshow_cls
|
||||
from utils.torch_utils import (ModelEMA, model_info, reshape_classifier_output, select_device, smart_DDP,
|
||||
smart_optimizer, smartCrossEntropyLoss, torch_distributed_zero_first)
|
||||
|
||||
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
|
||||
|
||||
|
||||
def train(opt, device):
|
||||
init_seeds(opt.seed + 1 + RANK, deterministic=True)
|
||||
save_dir, data, bs, epochs, nw, imgsz, pretrained = \
|
||||
opt.save_dir, Path(opt.data), opt.batch_size, opt.epochs, min(os.cpu_count() - 1, opt.workers), \
|
||||
opt.imgsz, str(opt.pretrained).lower() == 'true'
|
||||
cuda = device.type != 'cpu'
|
||||
|
||||
# Directories
|
||||
wdir = save_dir / 'weights'
|
||||
wdir.mkdir(parents=True, exist_ok=True) # make dir
|
||||
last, best = wdir / 'last.pt', wdir / 'best.pt'
|
||||
|
||||
# Save run settings
|
||||
yaml_save(save_dir / 'opt.yaml', vars(opt))
|
||||
|
||||
# Logger
|
||||
logger = GenericLogger(opt=opt, console_logger=LOGGER) if RANK in {-1, 0} else None
|
||||
|
||||
# Download Dataset
|
||||
with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
|
||||
data_dir = data if data.is_dir() else (DATASETS_DIR / data)
|
||||
if not data_dir.is_dir():
|
||||
LOGGER.info(f'\nDataset not found ⚠️, missing path {data_dir}, attempting download...')
|
||||
t = time.time()
|
||||
if str(data) == 'imagenet':
|
||||
subprocess.run(f"bash {ROOT / 'data/scripts/get_imagenet.sh'}", shell=True, check=True)
|
||||
else:
|
||||
url = f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{data}.zip'
|
||||
download(url, dir=data_dir.parent)
|
||||
s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n"
|
||||
LOGGER.info(s)
|
||||
|
||||
# Dataloaders
|
||||
nc = len([x for x in (data_dir / 'train').glob('*') if x.is_dir()]) # number of classes
|
||||
trainloader = create_classification_dataloader(path=data_dir / 'train',
|
||||
imgsz=imgsz,
|
||||
batch_size=bs // WORLD_SIZE,
|
||||
augment=True,
|
||||
cache=opt.cache,
|
||||
rank=LOCAL_RANK,
|
||||
workers=nw)
|
||||
|
||||
test_dir = data_dir / 'test' if (data_dir / 'test').exists() else data_dir / 'val' # data/test or data/val
|
||||
if RANK in {-1, 0}:
|
||||
testloader = create_classification_dataloader(path=test_dir,
|
||||
imgsz=imgsz,
|
||||
batch_size=bs // WORLD_SIZE * 2,
|
||||
augment=False,
|
||||
cache=opt.cache,
|
||||
rank=-1,
|
||||
workers=nw)
|
||||
|
||||
# Model
|
||||
with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
|
||||
if Path(opt.model).is_file() or opt.model.endswith('.pt'):
|
||||
model = attempt_load(opt.model, device='cpu', fuse=False)
|
||||
elif opt.model in torchvision.models.__dict__: # TorchVision models i.e. resnet50, efficientnet_b0
|
||||
model = torchvision.models.__dict__[opt.model](weights='IMAGENET1K_V1' if pretrained else None)
|
||||
else:
|
||||
m = hub.list('ultralytics/yolov5') # + hub.list('pytorch/vision') # models
|
||||
raise ModuleNotFoundError(f'--model {opt.model} not found. Available models are: \n' + '\n'.join(m))
|
||||
if isinstance(model, DetectionModel):
|
||||
LOGGER.warning("WARNING: pass YOLOv5 classifier model with '-cls' suffix, i.e. '--model yolov5s-cls.pt'")
|
||||
model = ClassificationModel(model=model, nc=nc, cutoff=opt.cutoff or 10) # convert to classification model
|
||||
reshape_classifier_output(model, nc) # update class count
|
||||
for p in model.parameters():
|
||||
p.requires_grad = True # for training
|
||||
for m in model.modules():
|
||||
if not pretrained and hasattr(m, 'reset_parameters'):
|
||||
m.reset_parameters()
|
||||
if isinstance(m, torch.nn.Dropout) and opt.dropout is not None:
|
||||
m.p = opt.dropout # set dropout
|
||||
model = model.to(device)
|
||||
names = trainloader.dataset.classes # class names
|
||||
model.names = names # attach class names
|
||||
|
||||
# Info
|
||||
if RANK in {-1, 0}:
|
||||
model_info(model)
|
||||
if opt.verbose:
|
||||
LOGGER.info(model)
|
||||
images, labels = next(iter(trainloader))
|
||||
file = imshow_cls(images[:25], labels[:25], names=names, f=save_dir / 'train_images.jpg')
|
||||
logger.log_images(file, name='Train Examples')
|
||||
logger.log_graph(model, imgsz) # log model
|
||||
|
||||
# Optimizer
|
||||
optimizer = smart_optimizer(model, opt.optimizer, opt.lr0, momentum=0.9, decay=5e-5)
|
||||
|
||||
# Scheduler
|
||||
lrf = 0.01 # final lr (fraction of lr0)
|
||||
# lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - lrf) + lrf # cosine
|
||||
lf = lambda x: (1 - x / epochs) * (1 - lrf) + lrf # linear
|
||||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
|
||||
# scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr0, total_steps=epochs, pct_start=0.1,
|
||||
# final_div_factor=1 / 25 / lrf)
|
||||
|
||||
# EMA
|
||||
ema = ModelEMA(model) if RANK in {-1, 0} else None
|
||||
|
||||
# DDP mode
|
||||
if cuda and RANK != -1:
|
||||
model = smart_DDP(model)
|
||||
|
||||
# Train
|
||||
t0 = time.time()
|
||||
criterion = smartCrossEntropyLoss(label_smoothing=opt.label_smoothing) # loss function
|
||||
best_fitness = 0.0
|
||||
scaler = amp.GradScaler(enabled=cuda)
|
||||
val = test_dir.stem # 'val' or 'test'
|
||||
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} test\n'
|
||||
f'Using {nw * WORLD_SIZE} dataloader workers\n'
|
||||
f"Logging results to {colorstr('bold', save_dir)}\n"
|
||||
f'Starting {opt.model} training on {data} dataset with {nc} classes for {epochs} epochs...\n\n'
|
||||
f"{'Epoch':>10}{'GPU_mem':>10}{'train_loss':>12}{f'{val}_loss':>12}{'top1_acc':>12}{'top5_acc':>12}")
|
||||
for epoch in range(epochs): # loop over the dataset multiple times
|
||||
tloss, vloss, fitness = 0.0, 0.0, 0.0 # train loss, val loss, fitness
|
||||
model.train()
|
||||
if RANK != -1:
|
||||
trainloader.sampler.set_epoch(epoch)
|
||||
pbar = enumerate(trainloader)
|
||||
if RANK in {-1, 0}:
|
||||
pbar = tqdm(enumerate(trainloader), total=len(trainloader), bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')
|
||||
for i, (images, labels) in pbar: # progress bar
|
||||
images, labels = images.to(device, non_blocking=True), labels.to(device)
|
||||
|
||||
# Forward
|
||||
with amp.autocast(enabled=cuda): # stability issues when enabled
|
||||
loss = criterion(model(images), labels)
|
||||
|
||||
# Backward
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
# Optimize
|
||||
scaler.unscale_(optimizer) # unscale gradients
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad()
|
||||
if ema:
|
||||
ema.update(model)
|
||||
|
||||
if RANK in {-1, 0}:
|
||||
# Print
|
||||
tloss = (tloss * i + loss.item()) / (i + 1) # update mean losses
|
||||
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB)
|
||||
pbar.desc = f"{f'{epoch + 1}/{epochs}':>10}{mem:>10}{tloss:>12.3g}" + ' ' * 36
|
||||
|
||||
# Test
|
||||
if i == len(pbar) - 1: # last batch
|
||||
top1, top5, vloss = validate.run(model=ema.ema,
|
||||
dataloader=testloader,
|
||||
criterion=criterion,
|
||||
pbar=pbar) # test accuracy, loss
|
||||
fitness = top1 # define fitness as top1 accuracy
|
||||
|
||||
# Scheduler
|
||||
scheduler.step()
|
||||
|
||||
# Log metrics
|
||||
if RANK in {-1, 0}:
|
||||
# Best fitness
|
||||
if fitness > best_fitness:
|
||||
best_fitness = fitness
|
||||
|
||||
# Log
|
||||
metrics = {
|
||||
"train/loss": tloss,
|
||||
f"{val}/loss": vloss,
|
||||
"metrics/accuracy_top1": top1,
|
||||
"metrics/accuracy_top5": top5,
|
||||
"lr/0": optimizer.param_groups[0]['lr']} # learning rate
|
||||
logger.log_metrics(metrics, epoch)
|
||||
|
||||
# Save model
|
||||
final_epoch = epoch + 1 == epochs
|
||||
if (not opt.nosave) or final_epoch:
|
||||
ckpt = {
|
||||
'epoch': epoch,
|
||||
'best_fitness': best_fitness,
|
||||
'model': deepcopy(ema.ema).half(), # deepcopy(de_parallel(model)).half(),
|
||||
'ema': None, # deepcopy(ema.ema).half(),
|
||||
'updates': ema.updates,
|
||||
'optimizer': None, # optimizer.state_dict(),
|
||||
'opt': vars(opt),
|
||||
'date': datetime.now().isoformat()}
|
||||
|
||||
# Save last, best and delete
|
||||
torch.save(ckpt, last)
|
||||
if best_fitness == fitness:
|
||||
torch.save(ckpt, best)
|
||||
del ckpt
|
||||
|
||||
# Train complete
|
||||
if RANK in {-1, 0} and final_epoch:
|
||||
LOGGER.info(f'\nTraining complete ({(time.time() - t0) / 3600:.3f} hours)'
|
||||
f"\nResults saved to {colorstr('bold', save_dir)}"
|
||||
f"\nPredict: python classify/predict.py --weights {best} --source im.jpg"
|
||||
f"\nValidate: python classify/val.py --weights {best} --data {data_dir}"
|
||||
f"\nExport: python export.py --weights {best} --include onnx"
|
||||
f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{best}')"
|
||||
f"\nVisualize: https://netron.app\n")
|
||||
|
||||
# Plot examples
|
||||
images, labels = (x[:25] for x in next(iter(testloader))) # first 25 images and labels
|
||||
pred = torch.max(ema.ema((images.half() if cuda else images.float()).to(device)), 1)[1]
|
||||
file = imshow_cls(images, labels, pred, names, verbose=False, f=save_dir / 'test_images.jpg')
|
||||
|
||||
# Log results
|
||||
meta = {"epochs": epochs, "top1_acc": best_fitness, "date": datetime.now().isoformat()}
|
||||
logger.log_images(file, name='Test Examples (true-predicted)', epoch=epoch)
|
||||
logger.log_model(best, epochs, metadata=meta)
|
||||
|
||||
|
||||
def parse_opt(known=False):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--model', type=str, default='yolov5s-cls.pt', help='initial weights path')
|
||||
parser.add_argument('--data', type=str, default='mnist', help='cifar10, cifar100, mnist, imagenet, etc.')
|
||||
parser.add_argument('--epochs', type=int, default=10)
|
||||
parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs')
|
||||
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=128, help='train, val image size (pixels)')
|
||||
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
||||
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
|
||||
parser.add_argument('--project', default=ROOT / 'runs/train-cls', 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('--pretrained', nargs='?', const=True, default=True, help='start from i.e. --pretrained False')
|
||||
parser.add_argument('--optimizer', choices=['SGD', 'Adam', 'AdamW', 'RMSProp'], default='Adam', help='optimizer')
|
||||
parser.add_argument('--lr0', type=float, default=0.001, help='initial learning rate')
|
||||
parser.add_argument('--label-smoothing', type=float, default=0.1, help='Label smoothing epsilon')
|
||||
parser.add_argument('--cutoff', type=int, default=None, help='Model layer cutoff index for Classify() head')
|
||||
parser.add_argument('--dropout', type=float, default=None, help='Dropout (fraction)')
|
||||
parser.add_argument('--verbose', action='store_true', help='Verbose mode')
|
||||
parser.add_argument('--seed', type=int, default=0, help='Global training seed')
|
||||
parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
|
||||
return parser.parse_known_args()[0] if known else parser.parse_args()
|
||||
|
||||
|
||||
def main(opt):
|
||||
# Checks
|
||||
if RANK in {-1, 0}:
|
||||
print_args(vars(opt))
|
||||
check_git_status()
|
||||
check_requirements()
|
||||
|
||||
# DDP mode
|
||||
device = select_device(opt.device, batch_size=opt.batch_size)
|
||||
if LOCAL_RANK != -1:
|
||||
assert opt.batch_size != -1, 'AutoBatch is coming soon for classification, please pass a valid --batch-size'
|
||||
assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
|
||||
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
|
||||
torch.cuda.set_device(LOCAL_RANK)
|
||||
device = torch.device('cuda', LOCAL_RANK)
|
||||
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
|
||||
|
||||
# Parameters
|
||||
opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
|
||||
|
||||
# Train
|
||||
train(opt, device)
|
||||
|
||||
|
||||
def run(**kwargs):
|
||||
# Usage: from yolov5 import classify; classify.train.run(data=mnist, imgsz=320, model='yolov5m')
|
||||
opt = parse_opt(True)
|
||||
for k, v in kwargs.items():
|
||||
setattr(opt, k, v)
|
||||
main(opt)
|
||||
return opt
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
opt = parse_opt()
|
||||
main(opt)
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Validate a classification model on a dataset
|
||||
|
||||
Usage:
|
||||
$ python classify/val.py --weights yolov5s-cls.pt --data ../datasets/imagenet
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[1] # YOLOv5 root directory
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||
|
||||
from models.common import DetectMultiBackend
|
||||
from utils.dataloaders import create_classification_dataloader
|
||||
from utils.general import LOGGER, check_img_size, check_requirements, colorstr, increment_path, print_args
|
||||
from utils.torch_utils import select_device, smart_inference_mode, time_sync
|
||||
|
||||
|
||||
@smart_inference_mode()
|
||||
def run(
|
||||
data=ROOT / '../datasets/mnist', # dataset dir
|
||||
weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s)
|
||||
batch_size=128, # batch size
|
||||
imgsz=224, # inference size (pixels)
|
||||
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||
workers=8, # max dataloader workers (per RANK in DDP mode)
|
||||
verbose=False, # verbose output
|
||||
project=ROOT / 'runs/val-cls', # save to project/name
|
||||
name='exp', # save to project/name
|
||||
exist_ok=False, # existing project/name ok, do not increment
|
||||
half=True, # use FP16 half-precision inference
|
||||
dnn=False, # use OpenCV DNN for ONNX inference
|
||||
model=None,
|
||||
dataloader=None,
|
||||
criterion=None,
|
||||
pbar=None,
|
||||
):
|
||||
# Initialize/load model and set device
|
||||
training = model is not None
|
||||
if training: # called by train.py
|
||||
device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model
|
||||
half &= device.type != 'cpu' # half precision only supported on CUDA
|
||||
model.half() if half else model.float()
|
||||
else: # called directly
|
||||
device = select_device(device, batch_size=batch_size)
|
||||
|
||||
# Directories
|
||||
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
||||
save_dir.mkdir(parents=True, exist_ok=True) # make dir
|
||||
|
||||
# Load model
|
||||
model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
|
||||
stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
|
||||
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||
half = model.fp16 # FP16 supported on limited backends with CUDA
|
||||
if engine:
|
||||
batch_size = model.batch_size
|
||||
else:
|
||||
device = model.device
|
||||
if not (pt or jit):
|
||||
batch_size = 1 # export.py models default to batch-size 1
|
||||
LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
|
||||
|
||||
# Dataloader
|
||||
data = Path(data)
|
||||
test_dir = data / 'test' if (data / 'test').exists() else data / 'val' # data/test or data/val
|
||||
dataloader = create_classification_dataloader(path=test_dir,
|
||||
imgsz=imgsz,
|
||||
batch_size=batch_size,
|
||||
augment=False,
|
||||
rank=-1,
|
||||
workers=workers)
|
||||
|
||||
model.eval()
|
||||
pred, targets, loss, dt = [], [], 0, [0.0, 0.0, 0.0]
|
||||
n = len(dataloader) # number of batches
|
||||
action = 'validating' if dataloader.dataset.root.stem == 'val' else 'testing'
|
||||
desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}"
|
||||
bar = tqdm(dataloader, desc, n, not training, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}', position=0)
|
||||
with torch.cuda.amp.autocast(enabled=device.type != 'cpu'):
|
||||
for images, labels in bar:
|
||||
t1 = time_sync()
|
||||
images, labels = images.to(device, non_blocking=True), labels.to(device)
|
||||
t2 = time_sync()
|
||||
dt[0] += t2 - t1
|
||||
|
||||
y = model(images)
|
||||
t3 = time_sync()
|
||||
dt[1] += t3 - t2
|
||||
|
||||
pred.append(y.argsort(1, descending=True)[:, :5])
|
||||
targets.append(labels)
|
||||
if criterion:
|
||||
loss += criterion(y, labels)
|
||||
dt[2] += time_sync() - t3
|
||||
|
||||
loss /= n
|
||||
pred, targets = torch.cat(pred), torch.cat(targets)
|
||||
correct = (targets[:, None] == pred).float()
|
||||
acc = torch.stack((correct[:, 0], correct.max(1).values), dim=1) # (top1, top5) accuracy
|
||||
top1, top5 = acc.mean(0).tolist()
|
||||
|
||||
if pbar:
|
||||
pbar.desc = f"{pbar.desc[:-36]}{loss:>12.3g}{top1:>12.3g}{top5:>12.3g}"
|
||||
if verbose: # all classes
|
||||
LOGGER.info(f"{'Class':>24}{'Images':>12}{'top1_acc':>12}{'top5_acc':>12}")
|
||||
LOGGER.info(f"{'all':>24}{targets.shape[0]:>12}{top1:>12.3g}{top5:>12.3g}")
|
||||
for i, c in enumerate(model.names):
|
||||
aci = acc[targets == i]
|
||||
top1i, top5i = aci.mean(0).tolist()
|
||||
LOGGER.info(f"{c:>24}{aci.shape[0]:>12}{top1i:>12.3g}{top5i:>12.3g}")
|
||||
|
||||
# Print results
|
||||
t = tuple(x / len(dataloader.dataset.samples) * 1E3 for x in dt) # speeds per image
|
||||
shape = (1, 3, imgsz, imgsz)
|
||||
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}' % t)
|
||||
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
|
||||
|
||||
return top1, top5, loss
|
||||
|
||||
|
||||
def parse_opt():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--data', type=str, default=ROOT / '../datasets/mnist', help='dataset path')
|
||||
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model.pt path(s)')
|
||||
parser.add_argument('--batch-size', type=int, default=128, help='batch size')
|
||||
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='inference size (pixels)')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
|
||||
parser.add_argument('--verbose', nargs='?', const=True, default=True, help='verbose output')
|
||||
parser.add_argument('--project', default=ROOT / 'runs/val-cls', 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')
|
||||
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
|
||||
opt = parser.parse_args()
|
||||
print_args(vars(opt))
|
||||
return opt
|
||||
|
||||
|
||||
def main(opt):
|
||||
check_requirements(exclude=('tensorboard', 'thop'))
|
||||
run(**vars(opt))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
opt = parse_opt()
|
||||
main(opt)
|
||||
Reference in new issue
Block a user