yolo with runs.zip file exp 14 is the best weight
This commit is contained in:
@@ -3,78 +3,16 @@
|
||||
utils/initialization
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import platform
|
||||
import threading
|
||||
|
||||
|
||||
def emojis(str=''):
|
||||
# Return platform-dependent emoji-safe version of string
|
||||
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
||||
|
||||
|
||||
class TryExcept(contextlib.ContextDecorator):
|
||||
# YOLOv3 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
|
||||
def __init__(self, msg=''):
|
||||
self.msg = msg
|
||||
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, exc_type, value, traceback):
|
||||
if value:
|
||||
print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
|
||||
return True
|
||||
|
||||
|
||||
def threaded(func):
|
||||
# Multi-threads a target function and returns thread. Usage: @threaded decorator
|
||||
def wrapper(*args, **kwargs):
|
||||
thread = threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True)
|
||||
thread.start()
|
||||
return thread
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def join_threads(verbose=False):
|
||||
# Join all daemon threads, i.e. atexit.register(lambda: join_threads())
|
||||
main_thread = threading.current_thread()
|
||||
for t in threading.enumerate():
|
||||
if t is not main_thread:
|
||||
if verbose:
|
||||
print(f'Joining thread {t.name}')
|
||||
t.join()
|
||||
|
||||
|
||||
def notebook_init(verbose=True):
|
||||
# Check system software and hardware
|
||||
def notebook_init():
|
||||
# For notebooks
|
||||
print('Checking setup...')
|
||||
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from utils.general import check_font, check_requirements, is_colab
|
||||
from utils.torch_utils import select_device # imports
|
||||
|
||||
check_font()
|
||||
|
||||
import psutil
|
||||
from IPython import display # to display images and clear console output
|
||||
|
||||
if is_colab():
|
||||
shutil.rmtree('/content/sample_data', ignore_errors=True) # remove colab /sample_data directory
|
||||
|
||||
# System info
|
||||
if verbose:
|
||||
gb = 1 << 30 # bytes to GiB (1024 ** 3)
|
||||
ram = psutil.virtual_memory().total
|
||||
total, used, free = shutil.disk_usage('/')
|
||||
display.clear_output()
|
||||
s = f'({os.cpu_count()} CPUs, {ram / gb:.1f} GB RAM, {(total - free) / gb:.1f}/{total / gb:.1f} GB disk)'
|
||||
else:
|
||||
s = ''
|
||||
from utils.general import emojis
|
||||
from utils.torch_utils import select_device # imports
|
||||
|
||||
display.clear_output()
|
||||
select_device(newline=False)
|
||||
print(emojis(f'Setup complete ✅ {s}'))
|
||||
print(emojis('Setup complete ✅'))
|
||||
return display
|
||||
|
||||
+10
-12
@@ -8,32 +8,29 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class SiLU(nn.Module):
|
||||
# SiLU activation https://arxiv.org/pdf/1606.08415.pdf
|
||||
# SiLU https://arxiv.org/pdf/1606.08415.pdf ----------------------------------------------------------------------------
|
||||
class SiLU(nn.Module): # export-friendly version of nn.SiLU()
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class Hardswish(nn.Module):
|
||||
# Hard-SiLU activation
|
||||
class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
# return x * F.hardsigmoid(x) # for TorchScript and CoreML
|
||||
return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for TorchScript, CoreML and ONNX
|
||||
# return x * F.hardsigmoid(x) # for torchscript and CoreML
|
||||
return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for torchscript, CoreML and ONNX
|
||||
|
||||
|
||||
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
||||
class Mish(nn.Module):
|
||||
# Mish activation https://github.com/digantamisra98/Mish
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x * F.softplus(x).tanh()
|
||||
|
||||
|
||||
class MemoryEfficientMish(nn.Module):
|
||||
# Mish activation memory-efficient
|
||||
class F(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x):
|
||||
ctx.save_for_backward(x)
|
||||
@@ -50,8 +47,8 @@ class MemoryEfficientMish(nn.Module):
|
||||
return self.F.apply(x)
|
||||
|
||||
|
||||
# FReLU https://arxiv.org/abs/2007.11824 -------------------------------------------------------------------------------
|
||||
class FReLU(nn.Module):
|
||||
# FReLU activation https://arxiv.org/abs/2007.11824
|
||||
def __init__(self, c1, k=3): # ch_in, kernel
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False)
|
||||
@@ -61,8 +58,9 @@ class FReLU(nn.Module):
|
||||
return torch.max(x, self.bn(self.conv(x)))
|
||||
|
||||
|
||||
# ACON https://arxiv.org/pdf/2009.04759.pdf ----------------------------------------------------------------------------
|
||||
class AconC(nn.Module):
|
||||
r""" ACON activation (activate or not)
|
||||
r""" ACON activation (activate or not).
|
||||
AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
|
||||
according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
|
||||
"""
|
||||
@@ -79,7 +77,7 @@ class AconC(nn.Module):
|
||||
|
||||
|
||||
class MetaAconC(nn.Module):
|
||||
r""" ACON activation (activate or not)
|
||||
r""" ACON activation (activate or not).
|
||||
MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network
|
||||
according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
|
||||
"""
|
||||
|
||||
+17
-137
@@ -8,42 +8,34 @@ import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms as T
|
||||
import torchvision.transforms.functional as TF
|
||||
|
||||
from utils.general import LOGGER, check_version, colorstr, resample_segments, segment2box, xywhn2xyxy
|
||||
from utils.general import LOGGER, check_version, colorstr, resample_segments, segment2box
|
||||
from utils.metrics import bbox_ioa
|
||||
|
||||
IMAGENET_MEAN = 0.485, 0.456, 0.406 # RGB mean
|
||||
IMAGENET_STD = 0.229, 0.224, 0.225 # RGB standard deviation
|
||||
|
||||
|
||||
class Albumentations:
|
||||
# YOLOv3 Albumentations class (optional, only used if package is installed)
|
||||
def __init__(self, size=640):
|
||||
# Albumentations class (optional, only used if package is installed)
|
||||
def __init__(self):
|
||||
self.transform = None
|
||||
prefix = colorstr('albumentations: ')
|
||||
try:
|
||||
import albumentations as A
|
||||
check_version(A.__version__, '1.0.3', hard=True) # version requirement
|
||||
|
||||
T = [
|
||||
A.RandomResizedCrop(height=size, width=size, scale=(0.8, 1.0), ratio=(0.9, 1.11), p=0.0),
|
||||
self.transform = A.Compose([
|
||||
A.Blur(p=0.01),
|
||||
A.MedianBlur(p=0.01),
|
||||
A.ToGray(p=0.01),
|
||||
A.CLAHE(p=0.01),
|
||||
A.RandomBrightnessContrast(p=0.0),
|
||||
A.RandomGamma(p=0.0),
|
||||
A.ImageCompression(quality_lower=75, p=0.0)] # transforms
|
||||
self.transform = A.Compose(T, bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
|
||||
A.ImageCompression(quality_lower=75, p=0.0)],
|
||||
bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
|
||||
|
||||
LOGGER.info(prefix + ', '.join(f'{x}'.replace('always_apply=False, ', '') for x in T if x.p))
|
||||
LOGGER.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms if x.p))
|
||||
except ImportError: # package not installed, skip
|
||||
pass
|
||||
except Exception as e:
|
||||
LOGGER.info(f'{prefix}{e}')
|
||||
LOGGER.info(colorstr('albumentations: ') + f'{e}')
|
||||
|
||||
def __call__(self, im, labels, p=1.0):
|
||||
if self.transform and random.random() < p:
|
||||
@@ -52,18 +44,6 @@ class Albumentations:
|
||||
return im, labels
|
||||
|
||||
|
||||
def normalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD, inplace=False):
|
||||
# Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = (x - mean) / std
|
||||
return TF.normalize(x, mean, std, inplace=inplace)
|
||||
|
||||
|
||||
def denormalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD):
|
||||
# Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = x * std + mean
|
||||
for i in range(3):
|
||||
x[:, i] = x[:, i] * std[i] + mean[i]
|
||||
return x
|
||||
|
||||
|
||||
def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
|
||||
# HSV color-space augmentation
|
||||
if hgain or sgain or vgain:
|
||||
@@ -141,14 +121,7 @@ def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleF
|
||||
return im, ratio, (dw, dh)
|
||||
|
||||
|
||||
def random_perspective(im,
|
||||
targets=(),
|
||||
segments=(),
|
||||
degrees=10,
|
||||
translate=.1,
|
||||
scale=.1,
|
||||
shear=10,
|
||||
perspective=0.0,
|
||||
def random_perspective(im, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0,
|
||||
border=(0, 0)):
|
||||
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10))
|
||||
# targets = [cls, xyxy]
|
||||
@@ -201,7 +174,7 @@ def random_perspective(im,
|
||||
# Transform label coordinates
|
||||
n = len(targets)
|
||||
if n:
|
||||
use_segments = any(x.any() for x in segments) and len(segments) == n
|
||||
use_segments = any(x.any() for x in segments)
|
||||
new = np.zeros((n, 4))
|
||||
if use_segments: # warp segments
|
||||
segments = resample_segments(segments) # upsample
|
||||
@@ -250,10 +223,12 @@ def copy_paste(im, labels, segments, p=0.5):
|
||||
if (ioa < 0.30).all(): # allow 30% obscuration of existing labels
|
||||
labels = np.concatenate((labels, [[l[0], *box]]), 0)
|
||||
segments.append(np.concatenate((w - s[:, 0:1], s[:, 1:2]), 1))
|
||||
cv2.drawContours(im_new, [segments[j].astype(np.int32)], -1, (1, 1, 1), cv2.FILLED)
|
||||
cv2.drawContours(im_new, [segments[j].astype(np.int32)], -1, (255, 255, 255), cv2.FILLED)
|
||||
|
||||
result = cv2.flip(im, 1) # augment segments (flip left-right)
|
||||
i = cv2.flip(im_new, 1).astype(bool)
|
||||
result = cv2.bitwise_and(src1=im, src2=im_new)
|
||||
result = cv2.flip(result, 1) # augment segments (flip left-right)
|
||||
i = result > 0 # pixels to replace
|
||||
# i[:, :] = result.max(2).reshape(h, w, 1) # act over ch
|
||||
im[i] = result[i] # cv2.imwrite('debug.jpg', im) # debug
|
||||
|
||||
return im, labels, segments
|
||||
@@ -280,7 +255,7 @@ def cutout(im, labels, p=0.5):
|
||||
# return unobscured labels
|
||||
if len(labels) and s > 0.03:
|
||||
box = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
|
||||
ioa = bbox_ioa(box, xywhn2xyxy(labels[:, 1:5], w, h)) # intersection over area
|
||||
ioa = bbox_ioa(box, labels[:, 1:5]) # intersection over area
|
||||
labels = labels[ioa < 0.60] # remove >60% obscured labels
|
||||
|
||||
return labels
|
||||
@@ -294,104 +269,9 @@ def mixup(im, labels, im2, labels2):
|
||||
return im, labels
|
||||
|
||||
|
||||
def box_candidates(box1, box2, wh_thr=2, ar_thr=100, area_thr=0.1, eps=1e-16): # box1(4,n), box2(4,n)
|
||||
def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1, eps=1e-16): # box1(4,n), box2(4,n)
|
||||
# Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio
|
||||
w1, h1 = box1[2] - box1[0], box1[3] - box1[1]
|
||||
w2, h2 = box2[2] - box2[0], box2[3] - box2[1]
|
||||
ar = np.maximum(w2 / (h2 + eps), h2 / (w2 + eps)) # aspect ratio
|
||||
return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr) # candidates
|
||||
|
||||
|
||||
def classify_albumentations(
|
||||
augment=True,
|
||||
size=224,
|
||||
scale=(0.08, 1.0),
|
||||
ratio=(0.75, 1.0 / 0.75), # 0.75, 1.33
|
||||
hflip=0.5,
|
||||
vflip=0.0,
|
||||
jitter=0.4,
|
||||
mean=IMAGENET_MEAN,
|
||||
std=IMAGENET_STD,
|
||||
auto_aug=False):
|
||||
# YOLOv3 classification Albumentations (optional, only used if package is installed)
|
||||
prefix = colorstr('albumentations: ')
|
||||
try:
|
||||
import albumentations as A
|
||||
from albumentations.pytorch import ToTensorV2
|
||||
check_version(A.__version__, '1.0.3', hard=True) # version requirement
|
||||
if augment: # Resize and crop
|
||||
T = [A.RandomResizedCrop(height=size, width=size, scale=scale, ratio=ratio)]
|
||||
if auto_aug:
|
||||
# TODO: implement AugMix, AutoAug & RandAug in albumentation
|
||||
LOGGER.info(f'{prefix}auto augmentations are currently not supported')
|
||||
else:
|
||||
if hflip > 0:
|
||||
T += [A.HorizontalFlip(p=hflip)]
|
||||
if vflip > 0:
|
||||
T += [A.VerticalFlip(p=vflip)]
|
||||
if jitter > 0:
|
||||
color_jitter = (float(jitter),) * 3 # repeat value for brightness, contrast, satuaration, 0 hue
|
||||
T += [A.ColorJitter(*color_jitter, 0)]
|
||||
else: # Use fixed crop for eval set (reproducibility)
|
||||
T = [A.SmallestMaxSize(max_size=size), A.CenterCrop(height=size, width=size)]
|
||||
T += [A.Normalize(mean=mean, std=std), ToTensorV2()] # Normalize and convert to Tensor
|
||||
LOGGER.info(prefix + ', '.join(f'{x}'.replace('always_apply=False, ', '') for x in T if x.p))
|
||||
return A.Compose(T)
|
||||
|
||||
except ImportError: # package not installed, skip
|
||||
LOGGER.warning(f'{prefix}⚠️ not found, install with `pip install albumentations` (recommended)')
|
||||
except Exception as e:
|
||||
LOGGER.info(f'{prefix}{e}')
|
||||
|
||||
|
||||
def classify_transforms(size=224):
|
||||
# Transforms to apply if albumentations not installed
|
||||
assert isinstance(size, int), f'ERROR: classify_transforms size {size} must be integer, not (list, tuple)'
|
||||
# T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
|
||||
return T.Compose([CenterCrop(size), ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
|
||||
|
||||
|
||||
class LetterBox:
|
||||
# YOLOv3 LetterBox class for image preprocessing, i.e. T.Compose([LetterBox(size), ToTensor()])
|
||||
def __init__(self, size=(640, 640), auto=False, stride=32):
|
||||
super().__init__()
|
||||
self.h, self.w = (size, size) if isinstance(size, int) else size
|
||||
self.auto = auto # pass max size integer, automatically solve for short side using stride
|
||||
self.stride = stride # used with auto
|
||||
|
||||
def __call__(self, im): # im = np.array HWC
|
||||
imh, imw = im.shape[:2]
|
||||
r = min(self.h / imh, self.w / imw) # ratio of new/old
|
||||
h, w = round(imh * r), round(imw * r) # resized image
|
||||
hs, ws = (math.ceil(x / self.stride) * self.stride for x in (h, w)) if self.auto else self.h, self.w
|
||||
top, left = round((hs - h) / 2 - 0.1), round((ws - w) / 2 - 0.1)
|
||||
im_out = np.full((self.h, self.w, 3), 114, dtype=im.dtype)
|
||||
im_out[top:top + h, left:left + w] = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
|
||||
return im_out
|
||||
|
||||
|
||||
class CenterCrop:
|
||||
# YOLOv3 CenterCrop class for image preprocessing, i.e. T.Compose([CenterCrop(size), ToTensor()])
|
||||
def __init__(self, size=640):
|
||||
super().__init__()
|
||||
self.h, self.w = (size, size) if isinstance(size, int) else size
|
||||
|
||||
def __call__(self, im): # im = np.array HWC
|
||||
imh, imw = im.shape[:2]
|
||||
m = min(imh, imw) # min dimension
|
||||
top, left = (imh - m) // 2, (imw - m) // 2
|
||||
return cv2.resize(im[top:top + m, left:left + m], (self.w, self.h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
|
||||
class ToTensor:
|
||||
# YOLOv3 ToTensor class for image preprocessing, i.e. T.Compose([LetterBox(size), ToTensor()])
|
||||
def __init__(self, half=False):
|
||||
super().__init__()
|
||||
self.half = half
|
||||
|
||||
def __call__(self, im): # im = np.array HWC in BGR order
|
||||
im = np.ascontiguousarray(im.transpose((2, 0, 1))[::-1]) # HWC to CHW -> BGR to RGB -> contiguous
|
||||
im = torch.from_numpy(im) # to torch
|
||||
im = im.half() if self.half else im.float() # uint8 to fp16/32
|
||||
im /= 255.0 # 0-255 to 0.0-1.0
|
||||
return im
|
||||
|
||||
+32
-37
@@ -1,6 +1,6 @@
|
||||
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
AutoAnchor utils
|
||||
Auto-anchor utils
|
||||
"""
|
||||
|
||||
import random
|
||||
@@ -10,23 +10,21 @@ import torch
|
||||
import yaml
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils import TryExcept
|
||||
from utils.general import LOGGER, TQDM_BAR_FORMAT, colorstr
|
||||
from utils.general import LOGGER, colorstr, emojis
|
||||
|
||||
PREFIX = colorstr('AutoAnchor: ')
|
||||
|
||||
|
||||
def check_anchor_order(m):
|
||||
# Check anchor order against stride order for YOLOv3 Detect() module m, and correct if necessary
|
||||
a = m.anchors.prod(-1).mean(-1).view(-1) # mean anchor area per output layer
|
||||
# Check anchor order against stride order for Detect() module m, and correct if necessary
|
||||
a = m.anchors.prod(-1).view(-1) # anchor area
|
||||
da = a[-1] - a[0] # delta a
|
||||
ds = m.stride[-1] - m.stride[0] # delta s
|
||||
if da and (da.sign() != ds.sign()): # same order
|
||||
if da.sign() != ds.sign(): # same order
|
||||
LOGGER.info(f'{PREFIX}Reversing anchor order')
|
||||
m.anchors[:] = m.anchors.flip(0)
|
||||
|
||||
|
||||
@TryExcept(f'{PREFIX}ERROR')
|
||||
def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
||||
# Check anchor fit to data, recompute if necessary
|
||||
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
|
||||
@@ -42,26 +40,26 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
||||
bpr = (best > 1 / thr).float().mean() # best possible recall
|
||||
return bpr, aat
|
||||
|
||||
stride = m.stride.to(m.anchors.device).view(-1, 1, 1) # model strides
|
||||
anchors = m.anchors.clone() * stride # current anchors
|
||||
anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors
|
||||
bpr, aat = metric(anchors.cpu().view(-1, 2))
|
||||
s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). '
|
||||
if bpr > 0.98: # threshold to recompute
|
||||
LOGGER.info(f'{s}Current anchors are a good fit to dataset ✅')
|
||||
LOGGER.info(emojis(f'{s}Current anchors are a good fit to dataset ✅'))
|
||||
else:
|
||||
LOGGER.info(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...')
|
||||
LOGGER.info(emojis(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...'))
|
||||
na = m.anchors.numel() // 2 # number of anchors
|
||||
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
||||
try:
|
||||
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
||||
except Exception as e:
|
||||
LOGGER.info(f'{PREFIX}ERROR: {e}')
|
||||
new_bpr = metric(anchors)[0]
|
||||
if new_bpr > bpr: # replace anchors
|
||||
anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
|
||||
m.anchors[:] = anchors.clone().view_as(m.anchors)
|
||||
check_anchor_order(m) # must be in pixel-space (not grid-space)
|
||||
m.anchors /= stride
|
||||
s = f'{PREFIX}Done ✅ (optional: update model *.yaml to use these anchors in the future)'
|
||||
m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
||||
check_anchor_order(m)
|
||||
LOGGER.info(f'{PREFIX}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
|
||||
else:
|
||||
s = f'{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)'
|
||||
LOGGER.info(s)
|
||||
LOGGER.info(f'{PREFIX}Original anchors better than new anchors. Proceeding with original anchors.')
|
||||
|
||||
|
||||
def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
|
||||
@@ -83,7 +81,6 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
|
||||
"""
|
||||
from scipy.cluster.vq import kmeans
|
||||
|
||||
npr = np.random
|
||||
thr = 1 / thr
|
||||
|
||||
def metric(k, wh): # compute metrics
|
||||
@@ -103,7 +100,7 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
|
||||
s = f'{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n' \
|
||||
f'{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' \
|
||||
f'past_thr={x[x > thr].mean():.3f}-mean: '
|
||||
for x in k:
|
||||
for i, x in enumerate(k):
|
||||
s += '%i,%i, ' % (round(x[0]), round(x[1]))
|
||||
if verbose:
|
||||
LOGGER.info(s[:-2])
|
||||
@@ -112,7 +109,7 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
|
||||
if isinstance(dataset, str): # *.yaml file
|
||||
with open(dataset, errors='ignore') as f:
|
||||
data_dict = yaml.safe_load(f) # model dict
|
||||
from utils.dataloaders import LoadImagesAndLabels
|
||||
from utils.datasets import LoadImagesAndLabels
|
||||
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
|
||||
|
||||
# Get label wh
|
||||
@@ -122,21 +119,18 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
|
||||
# Filter
|
||||
i = (wh0 < 3.0).any(1).sum()
|
||||
if i:
|
||||
LOGGER.info(f'{PREFIX}WARNING ⚠️ Extremely small objects found: {i} of {len(wh0)} labels are <3 pixels in size')
|
||||
wh = wh0[(wh0 >= 2.0).any(1)].astype(np.float32) # filter > 2 pixels
|
||||
# wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
|
||||
LOGGER.info(f'{PREFIX}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
|
||||
wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
|
||||
# wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
|
||||
|
||||
# Kmeans init
|
||||
try:
|
||||
LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
|
||||
assert n <= len(wh) # apply overdetermined constraint
|
||||
s = wh.std(0) # sigmas for whitening
|
||||
k = kmeans(wh / s, n, iter=30)[0] * s # points
|
||||
assert n == len(k) # kmeans may return fewer points than requested if wh is insufficient or too similar
|
||||
except Exception:
|
||||
LOGGER.warning(f'{PREFIX}WARNING ⚠️ switching strategies from kmeans to random init')
|
||||
k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size # random init
|
||||
wh, wh0 = (torch.tensor(x, dtype=torch.float32) for x in (wh, wh0))
|
||||
# Kmeans calculation
|
||||
LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
|
||||
s = wh.std(0) # sigmas for whitening
|
||||
k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
|
||||
assert len(k) == n, f'{PREFIX}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}'
|
||||
k *= s
|
||||
wh = torch.tensor(wh, dtype=torch.float32) # filtered
|
||||
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
|
||||
k = print_results(k, verbose=False)
|
||||
|
||||
# Plot
|
||||
@@ -152,8 +146,9 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
|
||||
# fig.savefig('wh.png', dpi=200)
|
||||
|
||||
# Evolve
|
||||
npr = np.random
|
||||
f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
|
||||
pbar = tqdm(range(gen), bar_format=TQDM_BAR_FORMAT) # progress bar
|
||||
pbar = tqdm(range(gen), desc=f'{PREFIX}Evolving anchors with Genetic Algorithm:') # progress bar
|
||||
for _ in pbar:
|
||||
v = np.ones(sh)
|
||||
while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
|
||||
@@ -166,4 +161,4 @@ def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen
|
||||
if verbose:
|
||||
print_results(k, verbose)
|
||||
|
||||
return print_results(k).astype(np.float32)
|
||||
return print_results(k)
|
||||
|
||||
+17
-32
@@ -7,66 +7,51 @@ from copy import deepcopy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.cuda import amp
|
||||
|
||||
from utils.general import LOGGER, colorstr
|
||||
from utils.torch_utils import profile
|
||||
|
||||
|
||||
def check_train_batch_size(model, imgsz=640, amp=True):
|
||||
# Check YOLOv3 training batch size
|
||||
with torch.cuda.amp.autocast(amp):
|
||||
def check_train_batch_size(model, imgsz=640):
|
||||
# Check training batch size
|
||||
with amp.autocast():
|
||||
return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size
|
||||
|
||||
|
||||
def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
|
||||
# Automatically estimate best YOLOv3 batch size to use `fraction` of available CUDA memory
|
||||
def autobatch(model, imgsz=640, fraction=0.9, batch_size=16):
|
||||
# Automatically estimate best batch size to use `fraction` of available CUDA memory
|
||||
# Usage:
|
||||
# import torch
|
||||
# from utils.autobatch import autobatch
|
||||
# model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)
|
||||
# model = torch.hub.load('ultralytics/yolov3', 'yolov3', autoshape=False)
|
||||
# print(autobatch(model))
|
||||
|
||||
# Check device
|
||||
prefix = colorstr('AutoBatch: ')
|
||||
LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}')
|
||||
device = next(model.parameters()).device # get model device
|
||||
if device.type == 'cpu':
|
||||
LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}')
|
||||
return batch_size
|
||||
if torch.backends.cudnn.benchmark:
|
||||
LOGGER.info(f'{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}')
|
||||
return batch_size
|
||||
|
||||
# Inspect CUDA memory
|
||||
gb = 1 << 30 # bytes to GiB (1024 ** 3)
|
||||
d = str(device).upper() # 'CUDA:0'
|
||||
properties = torch.cuda.get_device_properties(device) # device properties
|
||||
t = properties.total_memory / gb # GiB total
|
||||
r = torch.cuda.memory_reserved(device) / gb # GiB reserved
|
||||
a = torch.cuda.memory_allocated(device) / gb # GiB allocated
|
||||
f = t - (r + a) # GiB free
|
||||
t = properties.total_memory / 1024 ** 3 # (GiB)
|
||||
r = torch.cuda.memory_reserved(device) / 1024 ** 3 # (GiB)
|
||||
a = torch.cuda.memory_allocated(device) / 1024 ** 3 # (GiB)
|
||||
f = t - (r + a) # free inside reserved
|
||||
LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free')
|
||||
|
||||
# Profile batch sizes
|
||||
batch_sizes = [1, 2, 4, 8, 16]
|
||||
try:
|
||||
img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
|
||||
results = profile(img, model, n=3, device=device)
|
||||
img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes]
|
||||
y = profile(img, model, n=3, device=device)
|
||||
except Exception as e:
|
||||
LOGGER.warning(f'{prefix}{e}')
|
||||
|
||||
# Fit a solution
|
||||
y = [x[2] for x in results if x] # memory [2]
|
||||
p = np.polyfit(batch_sizes[:len(y)], y, deg=1) # first degree polynomial fit
|
||||
y = [x[2] for x in y if x] # memory [2]
|
||||
batch_sizes = batch_sizes[:len(y)]
|
||||
p = np.polyfit(batch_sizes, y, deg=1) # first degree polynomial fit
|
||||
b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size)
|
||||
if None in results: # some sizes failed
|
||||
i = results.index(None) # first fail index
|
||||
if b >= batch_sizes[i]: # y intercept above failure point
|
||||
b = batch_sizes[max(i - 1, 0)] # select prior safe point
|
||||
if b < 1 or b > 1024: # b outside of safe range
|
||||
b = batch_size
|
||||
LOGGER.warning(f'{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.')
|
||||
|
||||
fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted
|
||||
LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅')
|
||||
LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)')
|
||||
return b
|
||||
|
||||
+41
-41
@@ -3,46 +3,46 @@
|
||||
Callback utils
|
||||
"""
|
||||
|
||||
import threading
|
||||
|
||||
|
||||
class Callbacks:
|
||||
""""
|
||||
Handles all registered callbacks for YOLOv3 Hooks
|
||||
Handles all registered callbacks for Hooks
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Define the available callbacks
|
||||
self._callbacks = {
|
||||
'on_pretrain_routine_start': [],
|
||||
'on_pretrain_routine_end': [],
|
||||
'on_train_start': [],
|
||||
'on_train_epoch_start': [],
|
||||
'on_train_batch_start': [],
|
||||
'optimizer_step': [],
|
||||
'on_before_zero_grad': [],
|
||||
'on_train_batch_end': [],
|
||||
'on_train_epoch_end': [],
|
||||
'on_val_start': [],
|
||||
'on_val_batch_start': [],
|
||||
'on_val_image_end': [],
|
||||
'on_val_batch_end': [],
|
||||
'on_val_end': [],
|
||||
'on_fit_epoch_end': [], # fit = train + val
|
||||
'on_model_save': [],
|
||||
'on_train_end': [],
|
||||
'on_params_update': [],
|
||||
'teardown': [],}
|
||||
self.stop_training = False # set True to interrupt training
|
||||
# Define the available callbacks
|
||||
_callbacks = {
|
||||
'on_pretrain_routine_start': [],
|
||||
'on_pretrain_routine_end': [],
|
||||
|
||||
'on_train_start': [],
|
||||
'on_train_epoch_start': [],
|
||||
'on_train_batch_start': [],
|
||||
'optimizer_step': [],
|
||||
'on_before_zero_grad': [],
|
||||
'on_train_batch_end': [],
|
||||
'on_train_epoch_end': [],
|
||||
|
||||
'on_val_start': [],
|
||||
'on_val_batch_start': [],
|
||||
'on_val_image_end': [],
|
||||
'on_val_batch_end': [],
|
||||
'on_val_end': [],
|
||||
|
||||
'on_fit_epoch_end': [], # fit = train + val
|
||||
'on_model_save': [],
|
||||
'on_train_end': [],
|
||||
|
||||
'teardown': [],
|
||||
}
|
||||
|
||||
def register_action(self, hook, name='', callback=None):
|
||||
"""
|
||||
Register a new action to a callback hook
|
||||
|
||||
Args:
|
||||
hook: The callback hook name to register the action to
|
||||
name: The name of the action for later reference
|
||||
callback: The callback to fire
|
||||
hook The callback hook name to register the action to
|
||||
name The name of the action for later reference
|
||||
callback The callback to fire
|
||||
"""
|
||||
assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
|
||||
assert callable(callback), f"callback '{callback}' is not callable"
|
||||
@@ -53,24 +53,24 @@ class Callbacks:
|
||||
Returns all the registered actions by callback hook
|
||||
|
||||
Args:
|
||||
hook: The name of the hook to check, defaults to all
|
||||
hook The name of the hook to check, defaults to all
|
||||
"""
|
||||
return self._callbacks[hook] if hook else self._callbacks
|
||||
if hook:
|
||||
return self._callbacks[hook]
|
||||
else:
|
||||
return self._callbacks
|
||||
|
||||
def run(self, hook, *args, thread=False, **kwargs):
|
||||
def run(self, hook, *args, **kwargs):
|
||||
"""
|
||||
Loop through the registered actions and fire all callbacks on main thread
|
||||
Loop through the registered actions and fire all callbacks
|
||||
|
||||
Args:
|
||||
hook: The name of the hook to check, defaults to all
|
||||
args: Arguments to receive from YOLOv3
|
||||
thread: (boolean) Run callbacks in daemon thread
|
||||
kwargs: Keyword Arguments to receive from YOLOv3
|
||||
hook The name of the hook to check, defaults to all
|
||||
args Arguments to receive from
|
||||
kwargs Keyword Arguments to receive from
|
||||
"""
|
||||
|
||||
assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
|
||||
|
||||
for logger in self._callbacks[hook]:
|
||||
if thread:
|
||||
threading.Thread(target=logger['callback'], args=args, kwargs=kwargs, daemon=True).start()
|
||||
else:
|
||||
logger['callback'](*args, **kwargs)
|
||||
logger['callback'](*args, **kwargs)
|
||||
|
||||
+104
-61
@@ -3,104 +3,147 @@
|
||||
Download utils
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import time
|
||||
import urllib
|
||||
from pathlib import Path
|
||||
from zipfile import ZipFile
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
|
||||
def is_url(url, check=True):
|
||||
# Check if string is URL and check if URL exists
|
||||
try:
|
||||
url = str(url)
|
||||
result = urllib.parse.urlparse(url)
|
||||
assert all([result.scheme, result.netloc]) # check if is url
|
||||
return (urllib.request.urlopen(url).getcode() == 200) if check else True # check if exists online
|
||||
except (AssertionError, urllib.request.HTTPError):
|
||||
return False
|
||||
|
||||
|
||||
def gsutil_getsize(url=''):
|
||||
# gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
|
||||
s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8')
|
||||
return eval(s.split(' ')[0]) if len(s) else 0 # bytes
|
||||
|
||||
|
||||
def url_getsize(url='https://ultralytics.com/images/bus.jpg'):
|
||||
# Return downloadable file size in bytes
|
||||
response = requests.head(url, allow_redirects=True)
|
||||
return int(response.headers.get('content-length', -1))
|
||||
|
||||
|
||||
def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
|
||||
# Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
|
||||
from utils.general import LOGGER
|
||||
|
||||
file = Path(file)
|
||||
assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
|
||||
try: # url1
|
||||
LOGGER.info(f'Downloading {url} to {file}...')
|
||||
torch.hub.download_url_to_file(url, str(file), progress=LOGGER.level <= logging.INFO)
|
||||
print(f'Downloading {url} to {file}...')
|
||||
torch.hub.download_url_to_file(url, str(file))
|
||||
assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check
|
||||
except Exception as e: # url2
|
||||
if file.exists():
|
||||
file.unlink() # remove partial downloads
|
||||
LOGGER.info(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
|
||||
subprocess.run(
|
||||
f"curl -# -L '{url2 or url}' -o '{file}' --retry 3 -C -".split()) # curl download, retry and resume on fail
|
||||
file.unlink(missing_ok=True) # remove partial downloads
|
||||
print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
|
||||
os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail
|
||||
finally:
|
||||
if not file.exists() or file.stat().st_size < min_bytes: # check
|
||||
if file.exists():
|
||||
file.unlink() # remove partial downloads
|
||||
LOGGER.info(f'ERROR: {assert_msg}\n{error_msg}')
|
||||
LOGGER.info('')
|
||||
file.unlink(missing_ok=True) # remove partial downloads
|
||||
print(f"ERROR: {assert_msg}\n{error_msg}")
|
||||
print('')
|
||||
|
||||
|
||||
def attempt_download(file, repo='ultralytics/yolov5', release='v7.0'):
|
||||
# Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v7.0', etc.
|
||||
from utils.general import LOGGER
|
||||
|
||||
def github_assets(repository, version='latest'):
|
||||
# Return GitHub repo tag (i.e. 'v7.0') and assets (i.e. ['yolov5s.pt', 'yolov5m.pt', ...])
|
||||
if version != 'latest':
|
||||
version = f'tags/{version}' # i.e. tags/v7.0
|
||||
response = requests.get(f'https://api.github.com/repos/{repository}/releases/{version}').json() # github api
|
||||
return response['tag_name'], [x['name'] for x in response['assets']] # tag, assets
|
||||
|
||||
def attempt_download(file, repo='ultralytics/yolov3'): # from utils.downloads import *; attempt_download()
|
||||
# Attempt file download if does not exist
|
||||
file = Path(str(file).strip().replace("'", ''))
|
||||
|
||||
if not file.exists():
|
||||
# URL specified
|
||||
name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc.
|
||||
if str(file).startswith(('http:/', 'https:/')): # download
|
||||
url = str(file).replace(':/', '://') # Pathlib turns :// -> :/
|
||||
file = name.split('?')[0] # parse authentication https://url.com/file.txt?auth...
|
||||
if Path(file).is_file():
|
||||
LOGGER.info(f'Found {url} locally at {file}') # file already exists
|
||||
else:
|
||||
safe_download(file=file, url=url, min_bytes=1E5)
|
||||
return file
|
||||
name = name.split('?')[0] # parse authentication https://url.com/file.txt?auth...
|
||||
safe_download(file=name, url=url, min_bytes=1E5)
|
||||
return name
|
||||
|
||||
# GitHub assets
|
||||
assets = [f'yolov5{size}{suffix}.pt' for size in 'nsmlx' for suffix in ('', '6', '-cls', '-seg')] # default
|
||||
try:
|
||||
tag, assets = github_assets(repo, release)
|
||||
except Exception:
|
||||
try:
|
||||
tag, assets = github_assets(repo) # latest release
|
||||
except Exception:
|
||||
try:
|
||||
tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
|
||||
except Exception:
|
||||
tag = release
|
||||
|
||||
file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required)
|
||||
try:
|
||||
response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api
|
||||
assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov3.pt'...]
|
||||
tag = response['tag_name'] # i.e. 'v1.0'
|
||||
except: # fallback plan
|
||||
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
|
||||
try:
|
||||
tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
|
||||
except:
|
||||
tag = 'v9.5.0' # current release
|
||||
|
||||
if name in assets:
|
||||
safe_download(file,
|
||||
url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
|
||||
# url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
|
||||
min_bytes=1E5,
|
||||
error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/{tag}')
|
||||
error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
|
||||
|
||||
return str(file)
|
||||
|
||||
|
||||
def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
|
||||
# Downloads a file from Google Drive. from yolov3.utils.downloads import *; gdrive_download()
|
||||
t = time.time()
|
||||
file = Path(file)
|
||||
cookie = Path('cookie') # gdrive cookie
|
||||
print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='')
|
||||
file.unlink(missing_ok=True) # remove existing file
|
||||
cookie.unlink(missing_ok=True) # remove existing cookie
|
||||
|
||||
# Attempt file download
|
||||
out = "NUL" if platform.system() == "Windows" else "/dev/null"
|
||||
os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}')
|
||||
if os.path.exists('cookie'): # large file
|
||||
s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}'
|
||||
else: # small file
|
||||
s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"'
|
||||
r = os.system(s) # execute, capture return
|
||||
cookie.unlink(missing_ok=True) # remove existing cookie
|
||||
|
||||
# Error check
|
||||
if r != 0:
|
||||
file.unlink(missing_ok=True) # remove partial
|
||||
print('Download error ') # raise Exception('Download error')
|
||||
return r
|
||||
|
||||
# Unzip if archive
|
||||
if file.suffix == '.zip':
|
||||
print('unzipping... ', end='')
|
||||
ZipFile(file).extractall(path=file.parent) # unzip
|
||||
file.unlink() # remove zip
|
||||
|
||||
print(f'Done ({time.time() - t:.1f}s)')
|
||||
return r
|
||||
|
||||
|
||||
def get_token(cookie="./cookie"):
|
||||
with open(cookie) as f:
|
||||
for line in f:
|
||||
if "download" in line:
|
||||
return line.split()[-1]
|
||||
return ""
|
||||
|
||||
# Google utils: https://cloud.google.com/storage/docs/reference/libraries ----------------------------------------------
|
||||
#
|
||||
#
|
||||
# def upload_blob(bucket_name, source_file_name, destination_blob_name):
|
||||
# # Uploads a file to a bucket
|
||||
# # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
|
||||
#
|
||||
# storage_client = storage.Client()
|
||||
# bucket = storage_client.get_bucket(bucket_name)
|
||||
# blob = bucket.blob(destination_blob_name)
|
||||
#
|
||||
# blob.upload_from_filename(source_file_name)
|
||||
#
|
||||
# print('File {} uploaded to {}.'.format(
|
||||
# source_file_name,
|
||||
# destination_blob_name))
|
||||
#
|
||||
#
|
||||
# def download_blob(bucket_name, source_blob_name, destination_file_name):
|
||||
# # Uploads a blob from a bucket
|
||||
# storage_client = storage.Client()
|
||||
# bucket = storage_client.get_bucket(bucket_name)
|
||||
# blob = bucket.blob(source_blob_name)
|
||||
#
|
||||
# blob.download_to_filename(destination_file_name)
|
||||
#
|
||||
# print('Blob {} downloaded to {}.'.format(
|
||||
# source_blob_name,
|
||||
# destination_file_name))
|
||||
|
||||
Regular → Executable
+337
-631
File diff suppressed because it is too large
Load Diff
@@ -5,26 +5,25 @@ Logging utils
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
from threading import Thread
|
||||
|
||||
import pkg_resources as pkg
|
||||
import torch
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from utils.general import LOGGER, colorstr, cv2
|
||||
from utils.loggers.clearml.clearml_utils import ClearmlLogger
|
||||
from utils.general import colorstr, emojis
|
||||
from utils.loggers.wandb.wandb_utils import WandbLogger
|
||||
from utils.plots import plot_images, plot_labels, plot_results
|
||||
from utils.plots import plot_images, plot_results
|
||||
from utils.torch_utils import de_parallel
|
||||
|
||||
LOGGERS = ('csv', 'tb', 'wandb', 'clearml', 'comet') # *.csv, TensorBoard, Weights & Biases, ClearML
|
||||
LOGGERS = ('csv', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
|
||||
try:
|
||||
import wandb
|
||||
|
||||
assert hasattr(wandb, '__version__') # verify package import not local dir
|
||||
if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in {0, -1}:
|
||||
if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]:
|
||||
try:
|
||||
wandb_login_success = wandb.login(timeout=30)
|
||||
except wandb.errors.UsageError: # known non-TTY terminal issue
|
||||
@@ -34,64 +33,30 @@ try:
|
||||
except (ImportError, AssertionError):
|
||||
wandb = None
|
||||
|
||||
try:
|
||||
import clearml
|
||||
|
||||
assert hasattr(clearml, '__version__') # verify package import not local dir
|
||||
except (ImportError, AssertionError):
|
||||
clearml = None
|
||||
|
||||
try:
|
||||
if RANK not in [0, -1]:
|
||||
comet_ml = None
|
||||
else:
|
||||
import comet_ml
|
||||
|
||||
assert hasattr(comet_ml, '__version__') # verify package import not local dir
|
||||
from utils.loggers.comet import CometLogger
|
||||
|
||||
except (ModuleNotFoundError, ImportError, AssertionError):
|
||||
comet_ml = None
|
||||
|
||||
|
||||
class Loggers():
|
||||
# YOLOv3 Loggers class
|
||||
# Loggers class
|
||||
def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
|
||||
self.save_dir = save_dir
|
||||
self.weights = weights
|
||||
self.opt = opt
|
||||
self.hyp = hyp
|
||||
self.plots = not opt.noplots # plot results
|
||||
self.logger = logger # for printing results to console
|
||||
self.include = include
|
||||
self.keys = [
|
||||
'train/box_loss',
|
||||
'train/obj_loss',
|
||||
'train/cls_loss', # train loss
|
||||
'metrics/precision',
|
||||
'metrics/recall',
|
||||
'metrics/mAP_0.5',
|
||||
'metrics/mAP_0.5:0.95', # metrics
|
||||
'val/box_loss',
|
||||
'val/obj_loss',
|
||||
'val/cls_loss', # val loss
|
||||
'x/lr0',
|
||||
'x/lr1',
|
||||
'x/lr2'] # params
|
||||
self.best_keys = ['best/epoch', 'best/precision', 'best/recall', 'best/mAP_0.5', 'best/mAP_0.5:0.95']
|
||||
self.keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
|
||||
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', # metrics
|
||||
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
|
||||
'x/lr0', 'x/lr1', 'x/lr2'] # params
|
||||
for k in LOGGERS:
|
||||
setattr(self, k, None) # init empty logger dictionary
|
||||
self.csv = True # always log to csv
|
||||
|
||||
# Messages
|
||||
if not clearml:
|
||||
prefix = colorstr('ClearML: ')
|
||||
s = f"{prefix}run 'pip install clearml' to automatically track, visualize and remotely train YOLOv3 🚀 in ClearML"
|
||||
self.logger.info(s)
|
||||
if not comet_ml:
|
||||
prefix = colorstr('Comet: ')
|
||||
s = f"{prefix}run 'pip install comet_ml' to automatically track and visualize YOLOv3 🚀 runs in Comet"
|
||||
self.logger.info(s)
|
||||
# Message
|
||||
if not wandb:
|
||||
prefix = colorstr('Weights & Biases: ')
|
||||
s = f"{prefix}run 'pip install wandb' to automatically track and visualize YOLOv3 🚀 runs (RECOMMENDED)"
|
||||
print(emojis(s))
|
||||
|
||||
# TensorBoard
|
||||
s = self.save_dir
|
||||
if 'tb' in self.include and not self.opt.evolve:
|
||||
@@ -101,127 +66,53 @@ class Loggers():
|
||||
|
||||
# W&B
|
||||
if wandb and 'wandb' in self.include:
|
||||
wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith('wandb-artifact://')
|
||||
run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume and not wandb_artifact_resume else None
|
||||
self.opt.hyp = self.hyp # add hyperparameters
|
||||
self.wandb = WandbLogger(self.opt)
|
||||
self.wandb = WandbLogger(self.opt, run_id)
|
||||
else:
|
||||
self.wandb = None
|
||||
|
||||
# ClearML
|
||||
if clearml and 'clearml' in self.include:
|
||||
try:
|
||||
self.clearml = ClearmlLogger(self.opt, self.hyp)
|
||||
except Exception:
|
||||
self.clearml = None
|
||||
prefix = colorstr('ClearML: ')
|
||||
LOGGER.warning(f'{prefix}WARNING ⚠️ ClearML is installed but not configured, skipping ClearML logging.'
|
||||
f' See https://github.com/ultralytics/yolov5/tree/master/utils/loggers/clearml#readme')
|
||||
|
||||
else:
|
||||
self.clearml = None
|
||||
|
||||
# Comet
|
||||
if comet_ml and 'comet' in self.include:
|
||||
if isinstance(self.opt.resume, str) and self.opt.resume.startswith('comet://'):
|
||||
run_id = self.opt.resume.split('/')[-1]
|
||||
self.comet_logger = CometLogger(self.opt, self.hyp, run_id=run_id)
|
||||
|
||||
else:
|
||||
self.comet_logger = CometLogger(self.opt, self.hyp)
|
||||
|
||||
else:
|
||||
self.comet_logger = None
|
||||
|
||||
@property
|
||||
def remote_dataset(self):
|
||||
# Get data_dict if custom dataset artifact link is provided
|
||||
data_dict = None
|
||||
if self.clearml:
|
||||
data_dict = self.clearml.data_dict
|
||||
if self.wandb:
|
||||
data_dict = self.wandb.data_dict
|
||||
if self.comet_logger:
|
||||
data_dict = self.comet_logger.data_dict
|
||||
|
||||
return data_dict
|
||||
|
||||
def on_train_start(self):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_train_start()
|
||||
|
||||
def on_pretrain_routine_start(self):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_pretrain_routine_start()
|
||||
|
||||
def on_pretrain_routine_end(self, labels, names):
|
||||
def on_pretrain_routine_end(self):
|
||||
# Callback runs on pre-train routine end
|
||||
if self.plots:
|
||||
plot_labels(labels, names, self.save_dir)
|
||||
paths = self.save_dir.glob('*labels*.jpg') # training labels
|
||||
if self.wandb:
|
||||
self.wandb.log({'Labels': [wandb.Image(str(x), caption=x.name) for x in paths]})
|
||||
# if self.clearml:
|
||||
# pass # ClearML saves these images automatically using hooks
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_pretrain_routine_end(paths)
|
||||
paths = self.save_dir.glob('*labels*.jpg') # training labels
|
||||
if self.wandb:
|
||||
self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
|
||||
|
||||
def on_train_batch_end(self, model, ni, imgs, targets, paths, vals):
|
||||
log_dict = dict(zip(self.keys[:3], vals))
|
||||
def on_train_batch_end(self, ni, model, imgs, targets, paths, plots, sync_bn):
|
||||
# Callback runs on train batch end
|
||||
# ni: number integrated batches (since train start)
|
||||
if self.plots:
|
||||
if plots:
|
||||
if ni == 0:
|
||||
if not sync_bn: # tb.add_graph() --sync known issue https://github.com/ultralytics/yolov5/issues/3754
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('ignore') # suppress jit trace warning
|
||||
self.tb.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
|
||||
if ni < 3:
|
||||
f = self.save_dir / f'train_batch{ni}.jpg' # filename
|
||||
plot_images(imgs, targets, paths, f)
|
||||
if ni == 0 and self.tb and not self.opt.sync_bn:
|
||||
log_tensorboard_graph(self.tb, model, imgsz=(self.opt.imgsz, self.opt.imgsz))
|
||||
if ni == 10 and (self.wandb or self.clearml):
|
||||
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
|
||||
if self.wandb and ni == 10:
|
||||
files = sorted(self.save_dir.glob('train*.jpg'))
|
||||
if self.wandb:
|
||||
self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
|
||||
if self.clearml:
|
||||
self.clearml.log_debug_samples(files, title='Mosaics')
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_train_batch_end(log_dict, step=ni)
|
||||
self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
|
||||
|
||||
def on_train_epoch_end(self, epoch):
|
||||
# Callback runs on train epoch end
|
||||
if self.wandb:
|
||||
self.wandb.current_epoch = epoch + 1
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_train_epoch_end(epoch)
|
||||
|
||||
def on_val_start(self):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_val_start()
|
||||
|
||||
def on_val_image_end(self, pred, predn, path, names, im):
|
||||
# Callback runs on val image end
|
||||
if self.wandb:
|
||||
self.wandb.val_one_image(pred, predn, path, names, im)
|
||||
if self.clearml:
|
||||
self.clearml.log_image_with_boxes(path, pred, names, im)
|
||||
|
||||
def on_val_batch_end(self, batch_i, im, targets, paths, shapes, out):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_val_batch_end(batch_i, im, targets, paths, shapes, out)
|
||||
|
||||
def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
|
||||
def on_val_end(self):
|
||||
# Callback runs on val end
|
||||
if self.wandb or self.clearml:
|
||||
files = sorted(self.save_dir.glob('val*.jpg'))
|
||||
if self.wandb:
|
||||
self.wandb.log({'Validation': [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||
if self.clearml:
|
||||
self.clearml.log_debug_samples(files, title='Validation')
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_val_end(nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
|
||||
files = sorted(self.save_dir.glob('val*.jpg'))
|
||||
self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||
|
||||
def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
|
||||
# Callback runs at the end of each fit (train+val) epoch
|
||||
x = dict(zip(self.keys, vals))
|
||||
x = {k: v for k, v in zip(self.keys, vals)} # dict
|
||||
if self.csv:
|
||||
file = self.save_dir / 'results.csv'
|
||||
n = len(x) + 1 # number of cols
|
||||
@@ -232,170 +123,37 @@ class Loggers():
|
||||
if self.tb:
|
||||
for k, v in x.items():
|
||||
self.tb.add_scalar(k, v, epoch)
|
||||
elif self.clearml: # log to ClearML if TensorBoard not used
|
||||
for k, v in x.items():
|
||||
title, series = k.split('/')
|
||||
self.clearml.task.get_logger().report_scalar(title, series, v, epoch)
|
||||
|
||||
if self.wandb:
|
||||
if best_fitness == fi:
|
||||
best_results = [epoch] + vals[3:7]
|
||||
for i, name in enumerate(self.best_keys):
|
||||
self.wandb.wandb_run.summary[name] = best_results[i] # log best results in the summary
|
||||
self.wandb.log(x)
|
||||
self.wandb.end_epoch()
|
||||
|
||||
if self.clearml:
|
||||
self.clearml.current_epoch_logged_images = set() # reset epoch image limit
|
||||
self.clearml.current_epoch += 1
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_fit_epoch_end(x, epoch=epoch)
|
||||
self.wandb.end_epoch(best_result=best_fitness == fi)
|
||||
|
||||
def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
|
||||
# Callback runs on model save event
|
||||
if (epoch + 1) % self.opt.save_period == 0 and not final_epoch and self.opt.save_period != -1:
|
||||
if self.wandb:
|
||||
if self.wandb:
|
||||
if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
|
||||
self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
|
||||
if self.clearml:
|
||||
self.clearml.task.update_output_model(model_path=str(last),
|
||||
model_name='Latest Model',
|
||||
auto_delete_file=False)
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_model_save(last, epoch, final_epoch, best_fitness, fi)
|
||||
|
||||
def on_train_end(self, last, best, epoch, results):
|
||||
# Callback runs on training end, i.e. saving best model
|
||||
if self.plots:
|
||||
def on_train_end(self, last, best, plots, epoch, results):
|
||||
# Callback runs on training end
|
||||
if plots:
|
||||
plot_results(file=self.save_dir / 'results.csv') # save results.png
|
||||
files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
|
||||
files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter
|
||||
self.logger.info(f"Results saved to {colorstr('bold', self.save_dir)}")
|
||||
|
||||
if self.tb and not self.clearml: # These images are already captured by ClearML by now, we don't want doubles
|
||||
if self.tb:
|
||||
import cv2
|
||||
for f in files:
|
||||
self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
|
||||
|
||||
if self.wandb:
|
||||
self.wandb.log(dict(zip(self.keys[3:10], results)))
|
||||
self.wandb.log({'Results': [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||
self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||
# Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
|
||||
if not self.opt.evolve:
|
||||
wandb.log_artifact(str(best if best.exists() else last),
|
||||
type='model',
|
||||
name=f'run_{self.wandb.wandb_run.id}_model',
|
||||
wandb.log_artifact(str(best if best.exists() else last), type='model',
|
||||
name='run_' + self.wandb.wandb_run.id + '_model',
|
||||
aliases=['latest', 'best', 'stripped'])
|
||||
self.wandb.finish_run()
|
||||
|
||||
if self.clearml and not self.opt.evolve:
|
||||
self.clearml.task.update_output_model(model_path=str(best if best.exists() else last),
|
||||
name='Best Model',
|
||||
auto_delete_file=False)
|
||||
|
||||
if self.comet_logger:
|
||||
final_results = dict(zip(self.keys[3:10], results))
|
||||
self.comet_logger.on_train_end(files, self.save_dir, last, best, epoch, final_results)
|
||||
|
||||
def on_params_update(self, params: dict):
|
||||
# Update hyperparams or configs of the experiment
|
||||
if self.wandb:
|
||||
self.wandb.wandb_run.config.update(params, allow_val_change=True)
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_params_update(params)
|
||||
|
||||
|
||||
class GenericLogger:
|
||||
"""
|
||||
YOLOv5 General purpose logger for non-task specific logging
|
||||
Usage: from utils.loggers import GenericLogger; logger = GenericLogger(...)
|
||||
Arguments
|
||||
opt: Run arguments
|
||||
console_logger: Console logger
|
||||
include: loggers to include
|
||||
"""
|
||||
|
||||
def __init__(self, opt, console_logger, include=('tb', 'wandb')):
|
||||
# init default loggers
|
||||
self.save_dir = Path(opt.save_dir)
|
||||
self.include = include
|
||||
self.console_logger = console_logger
|
||||
self.csv = self.save_dir / 'results.csv' # CSV logger
|
||||
if 'tb' in self.include:
|
||||
prefix = colorstr('TensorBoard: ')
|
||||
self.console_logger.info(
|
||||
f"{prefix}Start with 'tensorboard --logdir {self.save_dir.parent}', view at http://localhost:6006/")
|
||||
self.tb = SummaryWriter(str(self.save_dir))
|
||||
|
||||
if wandb and 'wandb' in self.include:
|
||||
self.wandb = wandb.init(project=web_project_name(str(opt.project)),
|
||||
name=None if opt.name == 'exp' else opt.name,
|
||||
config=opt)
|
||||
else:
|
||||
self.wandb = None
|
||||
|
||||
def log_metrics(self, metrics, epoch):
|
||||
# Log metrics dictionary to all loggers
|
||||
if self.csv:
|
||||
keys, vals = list(metrics.keys()), list(metrics.values())
|
||||
n = len(metrics) + 1 # number of cols
|
||||
s = '' if self.csv.exists() else (('%23s,' * n % tuple(['epoch'] + keys)).rstrip(',') + '\n') # header
|
||||
with open(self.csv, 'a') as f:
|
||||
f.write(s + ('%23.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
|
||||
|
||||
if self.tb:
|
||||
for k, v in metrics.items():
|
||||
self.tb.add_scalar(k, v, epoch)
|
||||
|
||||
if self.wandb:
|
||||
self.wandb.log(metrics, step=epoch)
|
||||
|
||||
def log_images(self, files, name='Images', epoch=0):
|
||||
# Log images to all loggers
|
||||
files = [Path(f) for f in (files if isinstance(files, (tuple, list)) else [files])] # to Path
|
||||
files = [f for f in files if f.exists()] # filter by exists
|
||||
|
||||
if self.tb:
|
||||
for f in files:
|
||||
self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
|
||||
|
||||
if self.wandb:
|
||||
self.wandb.log({name: [wandb.Image(str(f), caption=f.name) for f in files]}, step=epoch)
|
||||
|
||||
def log_graph(self, model, imgsz=(640, 640)):
|
||||
# Log model graph to all loggers
|
||||
if self.tb:
|
||||
log_tensorboard_graph(self.tb, model, imgsz)
|
||||
|
||||
def log_model(self, model_path, epoch=0, metadata={}):
|
||||
# Log model to all loggers
|
||||
if self.wandb:
|
||||
art = wandb.Artifact(name=f'run_{wandb.run.id}_model', type='model', metadata=metadata)
|
||||
art.add_file(str(model_path))
|
||||
wandb.log_artifact(art)
|
||||
|
||||
def update_params(self, params):
|
||||
# Update the parameters logged
|
||||
if self.wandb:
|
||||
wandb.run.config.update(params, allow_val_change=True)
|
||||
|
||||
|
||||
def log_tensorboard_graph(tb, model, imgsz=(640, 640)):
|
||||
# Log model graph to TensorBoard
|
||||
try:
|
||||
p = next(model.parameters()) # for device, type
|
||||
imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz # expand
|
||||
im = torch.zeros((1, 3, *imgsz)).to(p.device).type_as(p) # input image (WARNING: must be zeros, not empty)
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('ignore') # suppress jit trace warning
|
||||
tb.add_graph(torch.jit.trace(de_parallel(model), im, strict=False), [])
|
||||
except Exception as e:
|
||||
LOGGER.warning(f'WARNING ⚠️ TensorBoard graph visualization failure {e}')
|
||||
|
||||
|
||||
def web_project_name(project):
|
||||
# Convert local project name to web project name
|
||||
if not project.startswith('runs/train'):
|
||||
return project
|
||||
suffix = '-Classify' if project.endswith('-cls') else '-Segment' if project.endswith('-seg') else ''
|
||||
return f'YOLOv5{suffix}'
|
||||
self.wandb.finish_run()
|
||||
else:
|
||||
self.wandb.finish_run()
|
||||
self.wandb = WandbLogger(self.opt)
|
||||
|
||||
@@ -1,32 +1,108 @@
|
||||
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||
|
||||
# WARNING ⚠️ wandb is deprecated and will be removed in future release.
|
||||
# See supported integrations at https://github.com/ultralytics/yolov5#integrations
|
||||
"""Utilities and tools for tracking runs with Weights & Biases."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
from utils.general import LOGGER, colorstr
|
||||
import pkg_resources as pkg
|
||||
import yaml
|
||||
from tqdm import tqdm
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[3] # YOLOv5 root directory
|
||||
ROOT = FILE.parents[3] # root directory
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
DEPRECATION_WARNING = f"{colorstr('wandb')}: WARNING ⚠️ wandb is deprecated and will be removed in a future release. " \
|
||||
f'See supported integrations at https://github.com/ultralytics/yolov5#integrations.'
|
||||
|
||||
from utils.datasets import LoadImagesAndLabels, img2label_paths
|
||||
from utils.general import LOGGER, check_dataset, check_file
|
||||
|
||||
try:
|
||||
import wandb
|
||||
|
||||
assert hasattr(wandb, '__version__') # verify package import not local dir
|
||||
LOGGER.warning(DEPRECATION_WARNING)
|
||||
except (ImportError, AssertionError):
|
||||
wandb = None
|
||||
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
||||
|
||||
|
||||
def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX):
|
||||
return from_string[len(prefix):]
|
||||
|
||||
|
||||
def check_wandb_config_file(data_config_file):
|
||||
wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path
|
||||
if Path(wandb_config).is_file():
|
||||
return wandb_config
|
||||
return data_config_file
|
||||
|
||||
|
||||
def check_wandb_dataset(data_file):
|
||||
is_trainset_wandb_artifact = False
|
||||
is_valset_wandb_artifact = False
|
||||
if check_file(data_file) and data_file.endswith('.yaml'):
|
||||
with open(data_file, errors='ignore') as f:
|
||||
data_dict = yaml.safe_load(f)
|
||||
is_trainset_wandb_artifact = (isinstance(data_dict['train'], str) and
|
||||
data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX))
|
||||
is_valset_wandb_artifact = (isinstance(data_dict['val'], str) and
|
||||
data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX))
|
||||
if is_trainset_wandb_artifact or is_valset_wandb_artifact:
|
||||
return data_dict
|
||||
else:
|
||||
return check_dataset(data_file)
|
||||
|
||||
|
||||
def get_run_info(run_path):
|
||||
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
|
||||
run_id = run_path.stem
|
||||
project = run_path.parent.stem
|
||||
entity = run_path.parent.parent.stem
|
||||
model_artifact_name = 'run_' + run_id + '_model'
|
||||
return entity, project, run_id, model_artifact_name
|
||||
|
||||
|
||||
def check_wandb_resume(opt):
|
||||
process_wandb_config_ddp_mode(opt) if RANK not in [-1, 0] else None
|
||||
if isinstance(opt.resume, str):
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
if RANK not in [-1, 0]: # For resuming DDP runs
|
||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||
api = wandb.Api()
|
||||
artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest')
|
||||
modeldir = artifact.download()
|
||||
opt.weights = str(Path(modeldir) / "last.pt")
|
||||
return True
|
||||
return None
|
||||
|
||||
|
||||
def process_wandb_config_ddp_mode(opt):
|
||||
with open(check_file(opt.data), errors='ignore') as f:
|
||||
data_dict = yaml.safe_load(f) # data dict
|
||||
train_dir, val_dir = None, None
|
||||
if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
|
||||
api = wandb.Api()
|
||||
train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias)
|
||||
train_dir = train_artifact.download()
|
||||
train_path = Path(train_dir) / 'data/images/'
|
||||
data_dict['train'] = str(train_path)
|
||||
|
||||
if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
|
||||
api = wandb.Api()
|
||||
val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
|
||||
val_dir = val_artifact.download()
|
||||
val_path = Path(val_dir) / 'data/images/'
|
||||
data_dict['val'] = str(val_path)
|
||||
if train_dir or val_dir:
|
||||
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
|
||||
with open(ddp_data_path, 'w') as f:
|
||||
yaml.safe_dump(data_dict, f)
|
||||
opt.data = ddp_data_path
|
||||
|
||||
|
||||
class WandbLogger():
|
||||
"""Log training runs, datasets, models, and predictions to Weights & Biases.
|
||||
@@ -46,7 +122,7 @@ class WandbLogger():
|
||||
"""
|
||||
- Initialize WandbLogger instance
|
||||
- Upload dataset if opt.upload_dataset is True
|
||||
- Setup training processes if job_type is 'Training'
|
||||
- Setup trainig processes if job_type is 'Training'
|
||||
|
||||
arguments:
|
||||
opt (namespace) -- Commandline arguments for this run
|
||||
@@ -56,31 +132,82 @@ class WandbLogger():
|
||||
"""
|
||||
# Pre-training routine --
|
||||
self.job_type = job_type
|
||||
self.wandb, self.wandb_run = wandb, wandb.run if wandb else None
|
||||
self.wandb, self.wandb_run = wandb, None if not wandb else wandb.run
|
||||
self.val_artifact, self.train_artifact = None, None
|
||||
self.train_artifact_path, self.val_artifact_path = None, None
|
||||
self.result_artifact = None
|
||||
self.val_table, self.result_table = None, None
|
||||
self.bbox_media_panel_images = []
|
||||
self.val_table_path_map = None
|
||||
self.max_imgs_to_log = 16
|
||||
self.wandb_artifact_data_dict = None
|
||||
self.data_dict = None
|
||||
if self.wandb:
|
||||
# It's more elegant to stick to 1 wandb.init call,
|
||||
# but useful config data is overwritten in the WandbLogger's wandb.init call
|
||||
if isinstance(opt.resume, str): # checks resume from artifact
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
|
||||
assert wandb, 'install wandb to resume wandb runs'
|
||||
# Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
|
||||
self.wandb_run = wandb.init(id=run_id,
|
||||
project=project,
|
||||
entity=entity,
|
||||
resume='allow',
|
||||
allow_val_change=True)
|
||||
opt.resume = model_artifact_name
|
||||
elif self.wandb:
|
||||
self.wandb_run = wandb.init(config=opt,
|
||||
resume='allow',
|
||||
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
|
||||
resume="allow",
|
||||
project='YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem,
|
||||
entity=opt.entity,
|
||||
name=opt.name if opt.name != 'exp' else None,
|
||||
job_type=job_type,
|
||||
id=run_id,
|
||||
allow_val_change=True) if not wandb.run else wandb.run
|
||||
|
||||
if self.wandb_run:
|
||||
if self.job_type == 'Training':
|
||||
if isinstance(opt.data, dict):
|
||||
# This means another dataset manager has already processed the dataset info (e.g. ClearML)
|
||||
# and they will have stored the already processed dict in opt.data
|
||||
self.data_dict = opt.data
|
||||
if opt.upload_dataset:
|
||||
if not opt.resume:
|
||||
self.wandb_artifact_data_dict = self.check_and_upload_dataset(opt)
|
||||
|
||||
if opt.resume:
|
||||
# resume from artifact
|
||||
if isinstance(opt.resume, str) and opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
self.data_dict = dict(self.wandb_run.config.data_dict)
|
||||
else: # local resume
|
||||
self.data_dict = check_wandb_dataset(opt.data)
|
||||
else:
|
||||
self.data_dict = check_wandb_dataset(opt.data)
|
||||
self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict
|
||||
|
||||
# write data_dict to config. useful for resuming from artifacts. Do this only when not resuming.
|
||||
self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict},
|
||||
allow_val_change=True)
|
||||
self.setup_training(opt)
|
||||
|
||||
if self.job_type == 'Dataset Creation':
|
||||
self.data_dict = self.check_and_upload_dataset(opt)
|
||||
|
||||
def check_and_upload_dataset(self, opt):
|
||||
"""
|
||||
Check if the dataset format is compatible and upload it as W&B artifact
|
||||
|
||||
arguments:
|
||||
opt (namespace)-- Commandline arguments for current run
|
||||
|
||||
returns:
|
||||
Updated dataset info dictionary where local dataset paths are replaced by WAND_ARFACT_PREFIX links.
|
||||
"""
|
||||
assert wandb, 'Install wandb to upload dataset'
|
||||
config_path = self.log_dataset_artifact(opt.data,
|
||||
opt.single_cls,
|
||||
'YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem)
|
||||
LOGGER.info(f"Created dataset config file {config_path}")
|
||||
with open(config_path, errors='ignore') as f:
|
||||
wandb_data_dict = yaml.safe_load(f)
|
||||
return wandb_data_dict
|
||||
|
||||
def setup_training(self, opt):
|
||||
"""
|
||||
Setup the necessary processes for training YOLO models:
|
||||
@@ -95,18 +222,77 @@ class WandbLogger():
|
||||
self.log_dict, self.current_epoch = {}, 0
|
||||
self.bbox_interval = opt.bbox_interval
|
||||
if isinstance(opt.resume, str):
|
||||
model_dir, _ = self.download_model_artifact(opt)
|
||||
if model_dir:
|
||||
self.weights = Path(model_dir) / 'last.pt'
|
||||
modeldir, _ = self.download_model_artifact(opt)
|
||||
if modeldir:
|
||||
self.weights = Path(modeldir) / "last.pt"
|
||||
config = self.wandb_run.config
|
||||
opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp, opt.imgsz = str(
|
||||
opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str(
|
||||
self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \
|
||||
config.hyp, config.imgsz
|
||||
config.hyp
|
||||
data_dict = self.data_dict
|
||||
if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download
|
||||
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'),
|
||||
opt.artifact_alias)
|
||||
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'),
|
||||
opt.artifact_alias)
|
||||
|
||||
if self.train_artifact_path is not None:
|
||||
train_path = Path(self.train_artifact_path) / 'data/images/'
|
||||
data_dict['train'] = str(train_path)
|
||||
if self.val_artifact_path is not None:
|
||||
val_path = Path(self.val_artifact_path) / 'data/images/'
|
||||
data_dict['val'] = str(val_path)
|
||||
|
||||
if self.val_artifact is not None:
|
||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||
self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"])
|
||||
self.val_table = self.val_artifact.get("val")
|
||||
if self.val_table_path_map is None:
|
||||
self.map_val_table_path()
|
||||
if opt.bbox_interval == -1:
|
||||
self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
|
||||
if opt.evolve or opt.noplots:
|
||||
self.bbox_interval = opt.bbox_interval = opt.epochs + 1 # disable bbox_interval
|
||||
train_from_artifact = self.train_artifact_path is not None and self.val_artifact_path is not None
|
||||
# Update the the data_dict to point to local artifacts dir
|
||||
if train_from_artifact:
|
||||
self.data_dict = data_dict
|
||||
|
||||
def download_dataset_artifact(self, path, alias):
|
||||
"""
|
||||
download the model checkpoint artifact if the path starts with WANDB_ARTIFACT_PREFIX
|
||||
|
||||
arguments:
|
||||
path -- path of the dataset to be used for training
|
||||
alias (str)-- alias of the artifact to be download/used for training
|
||||
|
||||
returns:
|
||||
(str, wandb.Artifact) -- path of the downladed dataset and it's corresponding artifact object if dataset
|
||||
is found otherwise returns (None, None)
|
||||
"""
|
||||
if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
|
||||
dataset_artifact = wandb.use_artifact(artifact_path.as_posix().replace("\\", "/"))
|
||||
assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
|
||||
datadir = dataset_artifact.download()
|
||||
return datadir, dataset_artifact
|
||||
return None, None
|
||||
|
||||
def download_model_artifact(self, opt):
|
||||
"""
|
||||
download the model checkpoint artifact if the resume path starts with WANDB_ARTIFACT_PREFIX
|
||||
|
||||
arguments:
|
||||
opt (namespace) -- Commandline arguments for this run
|
||||
"""
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest")
|
||||
assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
|
||||
modeldir = model_artifact.download()
|
||||
epochs_trained = model_artifact.metadata.get('epochs_trained')
|
||||
total_epochs = model_artifact.metadata.get('total_epochs')
|
||||
is_finished = total_epochs is None
|
||||
assert not is_finished, 'training is finished, can only resume incomplete runs.'
|
||||
return modeldir, model_artifact
|
||||
return None, None
|
||||
|
||||
def log_model(self, path, opt, epoch, fitness_score, best_model=False):
|
||||
"""
|
||||
@@ -119,22 +305,166 @@ class WandbLogger():
|
||||
fitness_score (float) -- fitness score for current epoch
|
||||
best_model (boolean) -- Boolean representing if the current checkpoint is the best yet.
|
||||
"""
|
||||
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model',
|
||||
type='model',
|
||||
metadata={
|
||||
'original_url': str(path),
|
||||
'epochs_trained': epoch + 1,
|
||||
'save period': opt.save_period,
|
||||
'project': opt.project,
|
||||
'total_epochs': opt.epochs,
|
||||
'fitness_score': fitness_score})
|
||||
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
|
||||
'original_url': str(path),
|
||||
'epochs_trained': epoch + 1,
|
||||
'save period': opt.save_period,
|
||||
'project': opt.project,
|
||||
'total_epochs': opt.epochs,
|
||||
'fitness_score': fitness_score
|
||||
})
|
||||
model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
|
||||
wandb.log_artifact(model_artifact,
|
||||
aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
|
||||
LOGGER.info(f'Saving model artifact on epoch {epoch + 1}')
|
||||
LOGGER.info(f"Saving model artifact on epoch {epoch + 1}")
|
||||
|
||||
def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
|
||||
"""
|
||||
Log the dataset as W&B artifact and return the new data file with W&B links
|
||||
|
||||
arguments:
|
||||
data_file (str) -- the .yaml file with information about the dataset like - path, classes etc.
|
||||
single_class (boolean) -- train multi-class data as single-class
|
||||
project (str) -- project name. Used to construct the artifact path
|
||||
overwrite_config (boolean) -- overwrites the data.yaml file if set to true otherwise creates a new
|
||||
file with _wandb postfix. Eg -> data_wandb.yaml
|
||||
|
||||
returns:
|
||||
the new .yaml file with artifact links. it can be used to start training directly from artifacts
|
||||
"""
|
||||
self.data_dict = check_dataset(data_file) # parse and check
|
||||
data = dict(self.data_dict)
|
||||
nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
|
||||
names = {k: v for k, v in enumerate(names)} # to index dictionary
|
||||
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None
|
||||
self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
|
||||
if data.get('train'):
|
||||
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
|
||||
if data.get('val'):
|
||||
data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')
|
||||
path = Path(data_file).stem
|
||||
path = (path if overwrite_config else path + '_wandb') + '.yaml' # updated data.yaml path
|
||||
data.pop('download', None)
|
||||
data.pop('path', None)
|
||||
with open(path, 'w') as f:
|
||||
yaml.safe_dump(data, f)
|
||||
|
||||
if self.job_type == 'Training': # builds correct artifact pipeline graph
|
||||
self.wandb_run.use_artifact(self.val_artifact)
|
||||
self.wandb_run.use_artifact(self.train_artifact)
|
||||
self.val_artifact.wait()
|
||||
self.val_table = self.val_artifact.get('val')
|
||||
self.map_val_table_path()
|
||||
else:
|
||||
self.wandb_run.log_artifact(self.train_artifact)
|
||||
self.wandb_run.log_artifact(self.val_artifact)
|
||||
return path
|
||||
|
||||
def map_val_table_path(self):
|
||||
"""
|
||||
Map the validation dataset Table like name of file -> it's id in the W&B Table.
|
||||
Useful for - referencing artifacts for evaluation.
|
||||
"""
|
||||
self.val_table_path_map = {}
|
||||
LOGGER.info("Mapping dataset")
|
||||
for i, data in enumerate(tqdm(self.val_table.data)):
|
||||
self.val_table_path_map[data[3]] = data[0]
|
||||
|
||||
def create_dataset_table(self, dataset: LoadImagesAndLabels, class_to_id: Dict[int,str], name: str = 'dataset'):
|
||||
"""
|
||||
Create and return W&B artifact containing W&B Table of the dataset.
|
||||
|
||||
arguments:
|
||||
dataset -- instance of LoadImagesAndLabels class used to iterate over the data to build Table
|
||||
class_to_id -- hash map that maps class ids to labels
|
||||
name -- name of the artifact
|
||||
|
||||
returns:
|
||||
dataset artifact to be logged or used
|
||||
"""
|
||||
# TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging
|
||||
artifact = wandb.Artifact(name=name, type="dataset")
|
||||
img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None
|
||||
img_files = tqdm(dataset.img_files) if not img_files else img_files
|
||||
for img_file in img_files:
|
||||
if Path(img_file).is_dir():
|
||||
artifact.add_dir(img_file, name='data/images')
|
||||
labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
|
||||
artifact.add_dir(labels_path, name='data/labels')
|
||||
else:
|
||||
artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
|
||||
label_file = Path(img2label_paths([img_file])[0])
|
||||
artifact.add_file(str(label_file),
|
||||
name='data/labels/' + label_file.name) if label_file.exists() else None
|
||||
table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
|
||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
|
||||
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
|
||||
box_data, img_classes = [], {}
|
||||
for cls, *xywh in labels[:, 1:].tolist():
|
||||
cls = int(cls)
|
||||
box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]},
|
||||
"class_id": cls,
|
||||
"box_caption": "%s" % (class_to_id[cls])})
|
||||
img_classes[cls] = class_to_id[cls]
|
||||
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
|
||||
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
|
||||
Path(paths).name)
|
||||
artifact.add(table, name)
|
||||
return artifact
|
||||
|
||||
def log_training_progress(self, predn, path, names):
|
||||
"""
|
||||
Build evaluation Table. Uses reference from validation dataset table.
|
||||
|
||||
arguments:
|
||||
predn (list): list of predictions in the native space in the format - [xmin, ymin, xmax, ymax, confidence, class]
|
||||
path (str): local path of the current evaluation image
|
||||
names (dict(int, str)): hash map that maps class ids to labels
|
||||
"""
|
||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()])
|
||||
box_data = []
|
||||
total_conf = 0
|
||||
for *xyxy, conf, cls in predn.tolist():
|
||||
if conf >= 0.25:
|
||||
box_data.append(
|
||||
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
"class_id": int(cls),
|
||||
"box_caption": f"{names[cls]} {conf:.3f}",
|
||||
"scores": {"class_score": conf},
|
||||
"domain": "pixel"})
|
||||
total_conf += conf
|
||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||
id = self.val_table_path_map[Path(path).name]
|
||||
self.result_table.add_data(self.current_epoch,
|
||||
id,
|
||||
self.val_table.data[id][1],
|
||||
wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
|
||||
total_conf / max(1, len(box_data))
|
||||
)
|
||||
|
||||
def val_one_image(self, pred, predn, path, names, im):
|
||||
pass
|
||||
"""
|
||||
Log validation data for one image. updates the result Table if validation dataset is uploaded and log bbox media panel
|
||||
|
||||
arguments:
|
||||
pred (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class]
|
||||
predn (list): list of predictions in the native space - [xmin, ymin, xmax, ymax, confidence, class]
|
||||
path (str): local path of the current evaluation image
|
||||
"""
|
||||
if self.val_table and self.result_table: # Log Table if Val dataset is uploaded as artifact
|
||||
self.log_training_progress(predn, path, names)
|
||||
|
||||
if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0:
|
||||
if self.current_epoch % self.bbox_interval == 0:
|
||||
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
"class_id": int(cls),
|
||||
"box_caption": f"{names[cls]} {conf:.3f}",
|
||||
"scores": {"class_score": conf},
|
||||
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||
self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name))
|
||||
|
||||
def log(self, log_dict):
|
||||
"""
|
||||
@@ -147,7 +477,7 @@ class WandbLogger():
|
||||
for key, value in log_dict.items():
|
||||
self.log_dict[key] = value
|
||||
|
||||
def end_epoch(self):
|
||||
def end_epoch(self, best_result=False):
|
||||
"""
|
||||
commit the log_dict, model artifacts and Tables to W&B and flush the log_dict.
|
||||
|
||||
@@ -156,15 +486,25 @@ class WandbLogger():
|
||||
"""
|
||||
if self.wandb_run:
|
||||
with all_logging_disabled():
|
||||
if self.bbox_media_panel_images:
|
||||
self.log_dict["BoundingBoxDebugger"] = self.bbox_media_panel_images
|
||||
try:
|
||||
wandb.log(self.log_dict)
|
||||
except BaseException as e:
|
||||
LOGGER.info(
|
||||
f'An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}'
|
||||
)
|
||||
LOGGER.info(f"An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}")
|
||||
self.wandb_run.finish()
|
||||
self.wandb_run = None
|
||||
|
||||
self.log_dict = {}
|
||||
self.bbox_media_panel_images = []
|
||||
if self.result_artifact:
|
||||
self.result_artifact.add(self.result_table, 'result')
|
||||
wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch),
|
||||
('best' if best_result else '')])
|
||||
|
||||
wandb.log({"evaluation": self.result_table})
|
||||
self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"])
|
||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||
|
||||
def finish_run(self):
|
||||
"""
|
||||
@@ -175,7 +515,6 @@ class WandbLogger():
|
||||
with all_logging_disabled():
|
||||
wandb.log(self.log_dict)
|
||||
wandb.run.finish()
|
||||
LOGGER.warning(DEPRECATION_WARNING)
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
||||
+86
-115
@@ -11,8 +11,6 @@ import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from utils import TryExcept, threaded
|
||||
|
||||
|
||||
def fitness(x):
|
||||
# Model fitness as a weighted combination of metrics
|
||||
@@ -20,15 +18,7 @@ def fitness(x):
|
||||
return (x[:, :4] * w).sum(1)
|
||||
|
||||
|
||||
def smooth(y, f=0.05):
|
||||
# Box filter of fraction f
|
||||
nf = round(len(y) * f * 2) // 2 + 1 # number of filter elements (must be odd)
|
||||
p = np.ones(nf // 2) # ones padding
|
||||
yp = np.concatenate((p * y[0], y, p * y[-1]), 0) # y padded
|
||||
return np.convolve(yp, np.ones(nf) / nf, mode='valid') # y-smoothed
|
||||
|
||||
|
||||
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16, prefix=''):
|
||||
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()):
|
||||
""" Compute the average precision, given the recall and precision curves.
|
||||
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
|
||||
# Arguments
|
||||
@@ -47,7 +37,7 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
|
||||
tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
|
||||
|
||||
# Find unique classes
|
||||
unique_classes, nt = np.unique(target_cls, return_counts=True)
|
||||
unique_classes = np.unique(target_cls)
|
||||
nc = unique_classes.shape[0] # number of classes, number of detections
|
||||
|
||||
# Create Precision-Recall curve and compute AP for each class
|
||||
@@ -55,44 +45,42 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
|
||||
ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
|
||||
for ci, c in enumerate(unique_classes):
|
||||
i = pred_cls == c
|
||||
n_l = nt[ci] # number of labels
|
||||
n_l = (target_cls == c).sum() # number of labels
|
||||
n_p = i.sum() # number of predictions
|
||||
|
||||
if n_p == 0 or n_l == 0:
|
||||
continue
|
||||
else:
|
||||
# Accumulate FPs and TPs
|
||||
fpc = (1 - tp[i]).cumsum(0)
|
||||
tpc = tp[i].cumsum(0)
|
||||
|
||||
# Accumulate FPs and TPs
|
||||
fpc = (1 - tp[i]).cumsum(0)
|
||||
tpc = tp[i].cumsum(0)
|
||||
# Recall
|
||||
recall = tpc / (n_l + 1e-16) # recall curve
|
||||
r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases
|
||||
|
||||
# Recall
|
||||
recall = tpc / (n_l + eps) # recall curve
|
||||
r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases
|
||||
# Precision
|
||||
precision = tpc / (tpc + fpc) # precision curve
|
||||
p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score
|
||||
|
||||
# Precision
|
||||
precision = tpc / (tpc + fpc) # precision curve
|
||||
p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score
|
||||
|
||||
# AP from recall-precision curve
|
||||
for j in range(tp.shape[1]):
|
||||
ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
|
||||
if plot and j == 0:
|
||||
py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
|
||||
# AP from recall-precision curve
|
||||
for j in range(tp.shape[1]):
|
||||
ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
|
||||
if plot and j == 0:
|
||||
py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
|
||||
|
||||
# Compute F1 (harmonic mean of precision and recall)
|
||||
f1 = 2 * p * r / (p + r + eps)
|
||||
f1 = 2 * p * r / (p + r + 1e-16)
|
||||
names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data
|
||||
names = dict(enumerate(names)) # to dict
|
||||
names = {i: v for i, v in enumerate(names)} # to dict
|
||||
if plot:
|
||||
plot_pr_curve(px, py, ap, Path(save_dir) / f'{prefix}PR_curve.png', names)
|
||||
plot_mc_curve(px, f1, Path(save_dir) / f'{prefix}F1_curve.png', names, ylabel='F1')
|
||||
plot_mc_curve(px, p, Path(save_dir) / f'{prefix}P_curve.png', names, ylabel='Precision')
|
||||
plot_mc_curve(px, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall')
|
||||
plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
|
||||
plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
|
||||
plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision')
|
||||
plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall')
|
||||
|
||||
i = smooth(f1.mean(0), 0.1).argmax() # max F1 index
|
||||
p, r, f1 = p[:, i], r[:, i], f1[:, i]
|
||||
tp = (r * nt).round() # true positives
|
||||
fp = (tp / (p + eps) - tp).round() # false positives
|
||||
return tp, fp, p, r, f1, ap, unique_classes.astype(int)
|
||||
i = f1.mean(0).argmax() # max F1 index
|
||||
return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32')
|
||||
|
||||
|
||||
def compute_ap(recall, precision):
|
||||
@@ -141,12 +129,6 @@ class ConfusionMatrix:
|
||||
Returns:
|
||||
None, updates confusion matrix accordingly
|
||||
"""
|
||||
if detections is None:
|
||||
gt_classes = labels.int()
|
||||
for gc in gt_classes:
|
||||
self.matrix[self.nc, gc] += 1 # background FN
|
||||
return
|
||||
|
||||
detections = detections[detections[:, 4] > self.conf]
|
||||
gt_classes = labels[:, 0].int()
|
||||
detection_classes = detections[:, 5].int()
|
||||
@@ -164,55 +146,43 @@ class ConfusionMatrix:
|
||||
matches = np.zeros((0, 3))
|
||||
|
||||
n = matches.shape[0] > 0
|
||||
m0, m1, _ = matches.transpose().astype(int)
|
||||
m0, m1, _ = matches.transpose().astype(np.int16)
|
||||
for i, gc in enumerate(gt_classes):
|
||||
j = m0 == i
|
||||
if n and sum(j) == 1:
|
||||
self.matrix[detection_classes[m1[j]], gc] += 1 # correct
|
||||
else:
|
||||
self.matrix[self.nc, gc] += 1 # true background
|
||||
self.matrix[self.nc, gc] += 1 # background FP
|
||||
|
||||
if n:
|
||||
for i, dc in enumerate(detection_classes):
|
||||
if not any(m1 == i):
|
||||
self.matrix[dc, self.nc] += 1 # predicted background
|
||||
self.matrix[dc, self.nc] += 1 # background FN
|
||||
|
||||
def tp_fp(self):
|
||||
tp = self.matrix.diagonal() # true positives
|
||||
fp = self.matrix.sum(1) - tp # false positives
|
||||
# fn = self.matrix.sum(0) - tp # false negatives (missed detections)
|
||||
return tp[:-1], fp[:-1] # remove background class
|
||||
def matrix(self):
|
||||
return self.matrix
|
||||
|
||||
@TryExcept('WARNING ⚠️ ConfusionMatrix plot failure')
|
||||
def plot(self, normalize=True, save_dir='', names=()):
|
||||
import seaborn as sn
|
||||
try:
|
||||
import seaborn as sn
|
||||
|
||||
array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-9) if normalize else 1) # normalize columns
|
||||
array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
|
||||
array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-6) if normalize else 1) # normalize columns
|
||||
array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
|
||||
|
||||
fig, ax = plt.subplots(1, 1, figsize=(12, 9), tight_layout=True)
|
||||
nc, nn = self.nc, len(names) # number of classes, names
|
||||
sn.set(font_scale=1.0 if nc < 50 else 0.8) # for label size
|
||||
labels = (0 < nn < 99) and (nn == nc) # apply names to ticklabels
|
||||
ticklabels = (names + ['background']) if labels else 'auto'
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered
|
||||
sn.heatmap(array,
|
||||
ax=ax,
|
||||
annot=nc < 30,
|
||||
annot_kws={
|
||||
'size': 8},
|
||||
cmap='Blues',
|
||||
fmt='.2f',
|
||||
square=True,
|
||||
vmin=0.0,
|
||||
xticklabels=ticklabels,
|
||||
yticklabels=ticklabels).set_facecolor((1, 1, 1))
|
||||
ax.set_xlabel('True')
|
||||
ax.set_ylabel('Predicted')
|
||||
ax.set_title('Confusion Matrix')
|
||||
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
|
||||
plt.close(fig)
|
||||
fig = plt.figure(figsize=(12, 9), tight_layout=True)
|
||||
sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size
|
||||
labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered
|
||||
sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
|
||||
xticklabels=names + ['background FP'] if labels else "auto",
|
||||
yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
|
||||
fig.axes[0].set_xlabel('True')
|
||||
fig.axes[0].set_ylabel('Predicted')
|
||||
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
|
||||
plt.close()
|
||||
except Exception as e:
|
||||
print(f'WARNING: ConfusionMatrix plot failure: {e}')
|
||||
|
||||
def print(self):
|
||||
for i in range(self.nc + 1):
|
||||
@@ -224,19 +194,19 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7
|
||||
|
||||
# Get the coordinates of bounding boxes
|
||||
if xywh: # transform from xywh to xyxy
|
||||
(x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
|
||||
(x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, 1), box2.chunk(4, 1)
|
||||
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_
|
||||
b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
|
||||
else: # x1, y1, x2, y2 = box1
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
|
||||
w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps)
|
||||
w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, 1)
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, 1)
|
||||
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
|
||||
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
|
||||
|
||||
# Intersection area
|
||||
inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * \
|
||||
(b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp(0)
|
||||
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
|
||||
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
|
||||
|
||||
# Union Area
|
||||
union = w1 * h1 + w2 * h2 - inter + eps
|
||||
@@ -244,13 +214,13 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7
|
||||
# IoU
|
||||
iou = inter / union
|
||||
if CIoU or DIoU or GIoU:
|
||||
cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1) # convex (smallest enclosing box) width
|
||||
ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1) # convex height
|
||||
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
|
||||
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
|
||||
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
|
||||
c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
|
||||
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: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
|
||||
v = (4 / math.pi ** 2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
|
||||
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
|
||||
with torch.no_grad():
|
||||
alpha = v / (v - iou + (1 + eps))
|
||||
return iou - (rho2 / c2 + v * alpha) # CIoU
|
||||
@@ -260,7 +230,7 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7
|
||||
return iou # IoU
|
||||
|
||||
|
||||
def box_iou(box1, box2, eps=1e-7):
|
||||
def box_iou(box1, box2):
|
||||
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
||||
"""
|
||||
Return intersection-over-union (Jaccard index) of boxes.
|
||||
@@ -273,24 +243,30 @@ def box_iou(box1, box2, eps=1e-7):
|
||||
IoU values for every element in boxes1 and boxes2
|
||||
"""
|
||||
|
||||
def box_area(box):
|
||||
# box = 4xn
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
area1 = box_area(box1.T)
|
||||
area2 = box_area(box2.T)
|
||||
|
||||
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
||||
(a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
|
||||
inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
|
||||
|
||||
# IoU = inter / (area1 + area2 - inter)
|
||||
return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
|
||||
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
||||
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
|
||||
|
||||
|
||||
def bbox_ioa(box1, box2, eps=1e-7):
|
||||
def bbox_ioa(box1, box2, eps=1E-7):
|
||||
""" Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2
|
||||
box1: np.array of shape(4)
|
||||
box2: np.array of shape(nx4)
|
||||
returns: np.array of shape(n)
|
||||
"""
|
||||
|
||||
box2 = box2.transpose()
|
||||
|
||||
# Get the coordinates of bounding boxes
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = box1
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = box2.T
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
|
||||
|
||||
# Intersection area
|
||||
inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \
|
||||
@@ -303,19 +279,17 @@ def bbox_ioa(box1, box2, eps=1e-7):
|
||||
return inter_area / box2_area
|
||||
|
||||
|
||||
def wh_iou(wh1, wh2, eps=1e-7):
|
||||
def wh_iou(wh1, wh2):
|
||||
# Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
|
||||
wh1 = wh1[:, None] # [N,1,2]
|
||||
wh2 = wh2[None] # [1,M,2]
|
||||
inter = torch.min(wh1, wh2).prod(2) # [N,M]
|
||||
return inter / (wh1.prod(2) + wh2.prod(2) - inter + eps) # iou = inter / (area1 + area2 - inter)
|
||||
return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
|
||||
|
||||
|
||||
# Plots ----------------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@threaded
|
||||
def plot_pr_curve(px, py, ap, save_dir=Path('pr_curve.png'), names=()):
|
||||
def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
|
||||
# Precision-recall curve
|
||||
fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
|
||||
py = np.stack(py, axis=1)
|
||||
@@ -331,14 +305,12 @@ def plot_pr_curve(px, py, ap, save_dir=Path('pr_curve.png'), names=()):
|
||||
ax.set_ylabel('Precision')
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left')
|
||||
ax.set_title('Precision-Recall Curve')
|
||||
fig.savefig(save_dir, dpi=250)
|
||||
plt.close(fig)
|
||||
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
|
||||
fig.savefig(Path(save_dir), dpi=250)
|
||||
plt.close()
|
||||
|
||||
|
||||
@threaded
|
||||
def plot_mc_curve(px, py, save_dir=Path('mc_curve.png'), names=(), xlabel='Confidence', ylabel='Metric'):
|
||||
def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'):
|
||||
# Metric-confidence curve
|
||||
fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
|
||||
|
||||
@@ -348,13 +320,12 @@ def plot_mc_curve(px, py, save_dir=Path('mc_curve.png'), names=(), xlabel='Confi
|
||||
else:
|
||||
ax.plot(px, py.T, linewidth=1, color='grey') # plot(confidence, metric)
|
||||
|
||||
y = smooth(py.mean(0), 0.05)
|
||||
y = py.mean(0)
|
||||
ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}')
|
||||
ax.set_xlabel(xlabel)
|
||||
ax.set_ylabel(ylabel)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.set_ylim(0, 1)
|
||||
ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left')
|
||||
ax.set_title(f'{ylabel}-Confidence Curve')
|
||||
fig.savefig(save_dir, dpi=250)
|
||||
plt.close(fig)
|
||||
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
|
||||
fig.savefig(Path(save_dir), dpi=250)
|
||||
plt.close()
|
||||
|
||||
+55
-146
@@ -3,12 +3,10 @@
|
||||
Plotting utils
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import math
|
||||
import os
|
||||
from copy import copy
|
||||
from pathlib import Path
|
||||
from urllib.error import URLError
|
||||
|
||||
import cv2
|
||||
import matplotlib
|
||||
@@ -19,13 +17,12 @@ import seaborn as sn
|
||||
import torch
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
from utils import TryExcept, threaded
|
||||
from utils.general import (CONFIG_DIR, FONT, LOGGER, check_font, check_requirements, clip_boxes, increment_path,
|
||||
is_ascii, xywh2xyxy, xyxy2xywh)
|
||||
from utils.general import (LOGGER, Timeout, check_requirements, clip_coords, increment_path, is_ascii, is_chinese,
|
||||
try_except, user_config_dir, xywh2xyxy, xyxy2xywh)
|
||||
from utils.metrics import fitness
|
||||
from utils.segment.general import scale_image
|
||||
|
||||
# Settings
|
||||
CONFIG_DIR = user_config_dir() # Ultralytics settings dir
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
matplotlib.rc('font', **{'size': 11})
|
||||
matplotlib.use('Agg') # for writing to files only
|
||||
@@ -35,9 +32,9 @@ class Colors:
|
||||
# Ultralytics color palette https://ultralytics.com/
|
||||
def __init__(self):
|
||||
# hex = matplotlib.colors.TABLEAU_COLORS.values()
|
||||
hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
|
||||
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
|
||||
self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
|
||||
hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
|
||||
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
|
||||
self.palette = [self.hex2rgb('#' + c) for c in hex]
|
||||
self.n = len(self.palette)
|
||||
|
||||
def __call__(self, i, bgr=False):
|
||||
@@ -52,33 +49,35 @@ class Colors:
|
||||
colors = Colors() # create instance for 'from utils.plots import colors'
|
||||
|
||||
|
||||
def check_pil_font(font=FONT, size=10):
|
||||
def check_font(font='Arial.ttf', size=10):
|
||||
# Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
|
||||
font = Path(font)
|
||||
font = font if font.exists() else (CONFIG_DIR / font.name)
|
||||
try:
|
||||
return ImageFont.truetype(str(font) if font.exists() else font.name, size)
|
||||
except Exception: # download if missing
|
||||
except Exception as e: # download if missing
|
||||
url = "https://ultralytics.com/assets/" + font.name
|
||||
print(f'Downloading {url} to {font}...')
|
||||
torch.hub.download_url_to_file(url, str(font), progress=False)
|
||||
try:
|
||||
check_font(font)
|
||||
return ImageFont.truetype(str(font), size)
|
||||
except TypeError:
|
||||
check_requirements('Pillow>=8.4.0') # known issue https://github.com/ultralytics/yolov5/issues/5374
|
||||
except URLError: # not online
|
||||
return ImageFont.load_default()
|
||||
|
||||
|
||||
class Annotator:
|
||||
# YOLOv3 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
||||
if RANK in (-1, 0):
|
||||
check_font() # download TTF if necessary
|
||||
|
||||
# Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
||||
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
||||
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
||||
non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
|
||||
self.pil = pil or non_ascii
|
||||
self.pil = pil or not is_ascii(example) or is_chinese(example)
|
||||
if self.pil: # use PIL
|
||||
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
||||
self.draw = ImageDraw.Draw(self.im)
|
||||
self.font = check_pil_font(font='Arial.Unicode.ttf' if non_ascii else font,
|
||||
size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
|
||||
self.font = check_font(font='Arial.Unicode.ttf' if is_chinese(example) else font,
|
||||
size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
|
||||
else: # use cv2
|
||||
self.im = im
|
||||
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
||||
@@ -88,14 +87,12 @@ class Annotator:
|
||||
if self.pil or not is_ascii(label):
|
||||
self.draw.rectangle(box, width=self.lw, outline=color) # box
|
||||
if label:
|
||||
w, h = self.font.getsize(label) # text width, height (WARNING: deprecated) in 9.2.0
|
||||
# _, _, w, h = self.font.getbbox(label) # text width, height (New)
|
||||
w, h = self.font.getsize(label) # text width, height
|
||||
outside = box[1] - h >= 0 # label fits outside box
|
||||
self.draw.rectangle(
|
||||
(box[0], box[1] - h if outside else box[1], box[0] + w + 1,
|
||||
box[1] + 1 if outside else box[1] + h + 1),
|
||||
fill=color,
|
||||
)
|
||||
self.draw.rectangle([box[0],
|
||||
box[1] - h if outside else box[1],
|
||||
box[0] + w + 1,
|
||||
box[1] + 1 if outside else box[1] + h + 1], fill=color)
|
||||
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
|
||||
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
|
||||
else: # cv2
|
||||
@@ -104,62 +101,20 @@ class Annotator:
|
||||
if label:
|
||||
tf = max(self.lw - 1, 1) # font thickness
|
||||
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
||||
outside = p1[1] - h >= 3
|
||||
outside = p1[1] - h - 3 >= 0 # label fits outside box
|
||||
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
||||
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
||||
cv2.putText(self.im,
|
||||
label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
|
||||
0,
|
||||
self.lw / 3,
|
||||
txt_color,
|
||||
thickness=tf,
|
||||
lineType=cv2.LINE_AA)
|
||||
|
||||
def masks(self, masks, colors, im_gpu, alpha=0.5, retina_masks=False):
|
||||
"""Plot masks at once.
|
||||
Args:
|
||||
masks (tensor): predicted masks on cuda, shape: [n, h, w]
|
||||
colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
|
||||
im_gpu (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
|
||||
alpha (float): mask transparency: 0.0 fully transparent, 1.0 opaque
|
||||
"""
|
||||
if self.pil:
|
||||
# convert to numpy first
|
||||
self.im = np.asarray(self.im).copy()
|
||||
if len(masks) == 0:
|
||||
self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
|
||||
colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
|
||||
colors = colors[:, None, None] # shape(n,1,1,3)
|
||||
masks = masks.unsqueeze(3) # shape(n,h,w,1)
|
||||
masks_color = masks * (colors * alpha) # shape(n,h,w,3)
|
||||
|
||||
inv_alph_masks = (1 - masks * alpha).cumprod(0) # shape(n,h,w,1)
|
||||
mcs = (masks_color * inv_alph_masks).sum(0) * 2 # mask color summand shape(n,h,w,3)
|
||||
|
||||
im_gpu = im_gpu.flip(dims=[0]) # flip channel
|
||||
im_gpu = im_gpu.permute(1, 2, 0).contiguous() # shape(h,w,3)
|
||||
im_gpu = im_gpu * inv_alph_masks[-1] + mcs
|
||||
im_mask = (im_gpu * 255).byte().cpu().numpy()
|
||||
self.im[:] = im_mask if retina_masks else scale_image(im_gpu.shape, im_mask, self.im.shape)
|
||||
if self.pil:
|
||||
# convert im back to PIL and update draw
|
||||
self.fromarray(self.im)
|
||||
cv2.putText(self.im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, self.lw / 3, txt_color,
|
||||
thickness=tf, lineType=cv2.LINE_AA)
|
||||
|
||||
def rectangle(self, xy, fill=None, outline=None, width=1):
|
||||
# Add rectangle to image (PIL-only)
|
||||
self.draw.rectangle(xy, fill, outline, width)
|
||||
|
||||
def text(self, xy, text, txt_color=(255, 255, 255), anchor='top'):
|
||||
def text(self, xy, text, txt_color=(255, 255, 255)):
|
||||
# Add text to image (PIL-only)
|
||||
if anchor == 'bottom': # start y from font bottom
|
||||
w, h = self.font.getsize(text) # text width, height
|
||||
xy[1] += 1 - h
|
||||
self.draw.text(xy, text, fill=txt_color, font=self.font)
|
||||
|
||||
def fromarray(self, im):
|
||||
# Update self.im from a numpy array
|
||||
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
||||
self.draw = ImageDraw.Draw(self.im)
|
||||
w, h = self.font.getsize(text) # text width, height
|
||||
self.draw.text((xy[0], xy[1] - h + 1), text, fill=txt_color, font=self.font)
|
||||
|
||||
def result(self):
|
||||
# Return annotated image as array
|
||||
@@ -177,7 +132,7 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
|
||||
if 'Detect' not in module_type:
|
||||
batch, channels, height, width = x.shape # batch, channels, height, width
|
||||
if height > 1 and width > 1:
|
||||
f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
|
||||
f = f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
|
||||
|
||||
blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels
|
||||
n = min(n, channels) # number of plots
|
||||
@@ -188,10 +143,9 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
|
||||
ax[i].imshow(blocks[i].squeeze()) # cmap='gray'
|
||||
ax[i].axis('off')
|
||||
|
||||
LOGGER.info(f'Saving {f}... ({n}/{channels})')
|
||||
plt.savefig(f, dpi=300, bbox_inches='tight')
|
||||
print(f'Saving {save_dir / f}... ({n}/{channels})')
|
||||
plt.savefig(save_dir / f, dpi=300, bbox_inches='tight')
|
||||
plt.close()
|
||||
np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy()) # npy save
|
||||
|
||||
|
||||
def hist2d(x, y, n=100):
|
||||
@@ -216,31 +170,26 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
||||
return filtfilt(b, a, data) # forward-backward filter
|
||||
|
||||
|
||||
def output_to_target(output, max_det=300):
|
||||
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting
|
||||
def output_to_target(output):
|
||||
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
|
||||
targets = []
|
||||
for i, o in enumerate(output):
|
||||
box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
|
||||
j = torch.full((conf.shape[0], 1), i)
|
||||
targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
|
||||
return torch.cat(targets, 0).numpy()
|
||||
for *box, conf, cls in o.cpu().numpy():
|
||||
targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf])
|
||||
return np.array(targets)
|
||||
|
||||
|
||||
@threaded
|
||||
def plot_images(images, targets, paths=None, fname='images.jpg', names=None):
|
||||
def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=1920, max_subplots=16):
|
||||
# Plot image grid with labels
|
||||
if isinstance(images, torch.Tensor):
|
||||
images = images.cpu().float().numpy()
|
||||
if isinstance(targets, torch.Tensor):
|
||||
targets = targets.cpu().numpy()
|
||||
|
||||
max_size = 1920 # max image size
|
||||
max_subplots = 16 # max image subplots, i.e. 4x4
|
||||
if np.max(images[0]) <= 1:
|
||||
images *= 255 # de-normalise (optional)
|
||||
bs, _, h, w = images.shape # batch size, _, height, width
|
||||
bs = min(bs, max_subplots) # limit plot images
|
||||
ns = np.ceil(bs ** 0.5) # number of subplots (square)
|
||||
if np.max(images[0]) <= 1:
|
||||
images *= 255 # de-normalise (optional)
|
||||
|
||||
# Build Image
|
||||
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
||||
@@ -260,12 +209,12 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None):
|
||||
|
||||
# Annotate
|
||||
fs = int((h + w) * ns * 0.01) # font size
|
||||
annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
|
||||
annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True)
|
||||
for i in range(i + 1):
|
||||
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
||||
annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders
|
||||
if paths:
|
||||
annotator.text((x + 5, y + 5), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
|
||||
annotator.text((x + 5, y + 5 + h), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
|
||||
if len(targets) > 0:
|
||||
ti = targets[targets[:, 0] == i] # image targets
|
||||
boxes = xywh2xyxy(ti[:, 2:6]).T
|
||||
@@ -346,7 +295,7 @@ def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_
|
||||
ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)[1].ravel()
|
||||
|
||||
fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
|
||||
# for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov5n6', 'yolov5s6', 'yolov5m6', 'yolov5l6', 'yolov5x6']]:
|
||||
# for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov3', 'yolov3-spp', 'yolov3-tiny']]:
|
||||
for f in sorted(save_dir.glob('study*.txt')):
|
||||
y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
|
||||
x = np.arange(y.shape[1]) if x is None else np.array(x)
|
||||
@@ -357,19 +306,11 @@ def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_
|
||||
ax[i].set_title(s[i])
|
||||
|
||||
j = y[3].argmax() + 1
|
||||
ax2.plot(y[5, 1:j],
|
||||
y[3, 1:j] * 1E2,
|
||||
'.-',
|
||||
linewidth=2,
|
||||
markersize=8,
|
||||
ax2.plot(y[5, 1:j], y[3, 1:j] * 1E2, '.-', linewidth=2, markersize=8,
|
||||
label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
|
||||
|
||||
ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
|
||||
'k.-',
|
||||
linewidth=2,
|
||||
markersize=8,
|
||||
alpha=.25,
|
||||
label='EfficientDet')
|
||||
'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
|
||||
|
||||
ax2.grid(alpha=0.2)
|
||||
ax2.set_yticks(np.arange(20, 60, 5))
|
||||
@@ -383,7 +324,8 @@ def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_
|
||||
plt.savefig(f, dpi=300)
|
||||
|
||||
|
||||
@TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395
|
||||
@try_except # known issue https://github.com/ultralytics/yolov5/issues/5395
|
||||
@Timeout(30) # known issue https://github.com/ultralytics/yolov5/issues/5611
|
||||
def plot_labels(labels, names=(), save_dir=Path('')):
|
||||
# plot dataset labels
|
||||
LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
|
||||
@@ -400,12 +342,11 @@ def plot_labels(labels, names=(), save_dir=Path('')):
|
||||
matplotlib.use('svg') # faster
|
||||
ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
|
||||
y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
|
||||
with contextlib.suppress(Exception): # color histogram bars by class
|
||||
[y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # known issue #3195
|
||||
# [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # update colors bug #3195
|
||||
ax[0].set_ylabel('instances')
|
||||
if 0 < len(names) < 30:
|
||||
ax[0].set_xticks(range(len(names)))
|
||||
ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10)
|
||||
ax[0].set_xticklabels(names, rotation=90, fontsize=10)
|
||||
else:
|
||||
ax[0].set_xlabel('classes')
|
||||
sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
|
||||
@@ -429,35 +370,6 @@ def plot_labels(labels, names=(), save_dir=Path('')):
|
||||
plt.close()
|
||||
|
||||
|
||||
def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f=Path('images.jpg')):
|
||||
# Show classification image grid with labels (optional) and predictions (optional)
|
||||
from utils.augmentations import denormalize
|
||||
|
||||
names = names or [f'class{i}' for i in range(1000)]
|
||||
blocks = torch.chunk(denormalize(im.clone()).cpu().float(), len(im),
|
||||
dim=0) # select batch index 0, block by channels
|
||||
n = min(len(blocks), nmax) # number of plots
|
||||
m = min(8, round(n ** 0.5)) # 8 x 8 default
|
||||
fig, ax = plt.subplots(math.ceil(n / m), m) # 8 rows x n/8 cols
|
||||
ax = ax.ravel() if m > 1 else [ax]
|
||||
# plt.subplots_adjust(wspace=0.05, hspace=0.05)
|
||||
for i in range(n):
|
||||
ax[i].imshow(blocks[i].squeeze().permute((1, 2, 0)).numpy().clip(0.0, 1.0))
|
||||
ax[i].axis('off')
|
||||
if labels is not None:
|
||||
s = names[labels[i]] + (f'—{names[pred[i]]}' if pred is not None else '')
|
||||
ax[i].set_title(s, fontsize=8, verticalalignment='top')
|
||||
plt.savefig(f, dpi=300, bbox_inches='tight')
|
||||
plt.close()
|
||||
if verbose:
|
||||
LOGGER.info(f'Saving {f}')
|
||||
if labels is not None:
|
||||
LOGGER.info('True: ' + ' '.join(f'{names[i]:3s}' for i in labels[:nmax]))
|
||||
if pred is not None:
|
||||
LOGGER.info('Predicted:' + ' '.join(f'{names[i]:3s}' for i in pred[:nmax]))
|
||||
return f
|
||||
|
||||
|
||||
def plot_evolve(evolve_csv='path/to/evolve.csv'): # from utils.plots import *; plot_evolve()
|
||||
# Plot evolve.csv hyp evolution results
|
||||
evolve_csv = Path(evolve_csv)
|
||||
@@ -468,7 +380,6 @@ def plot_evolve(evolve_csv='path/to/evolve.csv'): # from utils.plots import *;
|
||||
j = np.argmax(f) # max fitness index
|
||||
plt.figure(figsize=(10, 12), tight_layout=True)
|
||||
matplotlib.rc('font', **{'size': 8})
|
||||
print(f'Best results from row {j} of {evolve_csv}:')
|
||||
for i, k in enumerate(keys[7:]):
|
||||
v = x[:, 7 + i]
|
||||
mu = v[j] # best single result
|
||||
@@ -492,20 +403,20 @@ def plot_results(file='path/to/results.csv', dir=''):
|
||||
ax = ax.ravel()
|
||||
files = list(save_dir.glob('results*.csv'))
|
||||
assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.'
|
||||
for f in files:
|
||||
for fi, f in enumerate(files):
|
||||
try:
|
||||
data = pd.read_csv(f)
|
||||
s = [x.strip() for x in data.columns]
|
||||
x = data.values[:, 0]
|
||||
for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]):
|
||||
y = data.values[:, j].astype('float')
|
||||
y = data.values[:, j]
|
||||
# y[y == 0] = np.nan # don't show zero values
|
||||
ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=8)
|
||||
ax[i].set_title(s[j], fontsize=12)
|
||||
# if j in [8, 9, 10]: # share train and val loss y axes
|
||||
# ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
|
||||
except Exception as e:
|
||||
LOGGER.info(f'Warning: Plotting error for {f}: {e}')
|
||||
print(f'Warning: Plotting error for {f}: {e}')
|
||||
ax[1].legend()
|
||||
fig.savefig(save_dir / 'results.png', dpi=200)
|
||||
plt.close()
|
||||
@@ -542,7 +453,7 @@ def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
|
||||
plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
|
||||
|
||||
|
||||
def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
|
||||
def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
|
||||
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
|
||||
xyxy = torch.tensor(xyxy).view(-1, 4)
|
||||
b = xyxy2xywh(xyxy) # boxes
|
||||
@@ -550,11 +461,9 @@ def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False,
|
||||
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
||||
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
|
||||
xyxy = xywh2xyxy(b).long()
|
||||
clip_boxes(xyxy, im.shape)
|
||||
clip_coords(xyxy, im.shape)
|
||||
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
|
||||
if save:
|
||||
file.parent.mkdir(parents=True, exist_ok=True) # make directory
|
||||
f = str(increment_path(file).with_suffix('.jpg'))
|
||||
# cv2.imwrite(f, crop) # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
|
||||
Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0) # save RGB
|
||||
cv2.imwrite(str(increment_path(file).with_suffix('.jpg')), crop)
|
||||
return crop
|
||||
|
||||
+89
-203
@@ -3,12 +3,12 @@
|
||||
PyTorch utils
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import math
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import time
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
@@ -17,77 +17,20 @@ import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
from utils.general import LOGGER, check_version, colorstr, file_date, git_describe
|
||||
|
||||
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))
|
||||
from utils.general import LOGGER
|
||||
|
||||
try:
|
||||
import thop # for FLOPs computation
|
||||
except ImportError:
|
||||
thop = None
|
||||
|
||||
# Suppress PyTorch warnings
|
||||
warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')
|
||||
warnings.filterwarnings('ignore', category=UserWarning)
|
||||
|
||||
|
||||
def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
|
||||
# Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator
|
||||
def decorate(fn):
|
||||
return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
|
||||
|
||||
return decorate
|
||||
|
||||
|
||||
def smartCrossEntropyLoss(label_smoothing=0.0):
|
||||
# Returns nn.CrossEntropyLoss with label smoothing enabled for torch>=1.10.0
|
||||
if check_version(torch.__version__, '1.10.0'):
|
||||
return nn.CrossEntropyLoss(label_smoothing=label_smoothing)
|
||||
if label_smoothing > 0:
|
||||
LOGGER.warning(f'WARNING ⚠️ label smoothing {label_smoothing} requires torch>=1.10.0')
|
||||
return nn.CrossEntropyLoss()
|
||||
|
||||
|
||||
def smart_DDP(model):
|
||||
# Model DDP creation with checks
|
||||
assert not check_version(torch.__version__, '1.12.0', pinned=True), \
|
||||
'torch==1.12.0 torchvision==0.13.0 DDP training is not supported due to a known issue. ' \
|
||||
'Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395'
|
||||
if check_version(torch.__version__, '1.11.0'):
|
||||
return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK, static_graph=True)
|
||||
else:
|
||||
return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
|
||||
|
||||
|
||||
def reshape_classifier_output(model, n=1000):
|
||||
# Update a TorchVision classification model to class count 'n' if required
|
||||
from models.common import Classify
|
||||
name, m = list((model.model if hasattr(model, 'model') else model).named_children())[-1] # last module
|
||||
if isinstance(m, Classify): # YOLOv3 Classify() head
|
||||
if m.linear.out_features != n:
|
||||
m.linear = nn.Linear(m.linear.in_features, n)
|
||||
elif isinstance(m, nn.Linear): # ResNet, EfficientNet
|
||||
if m.out_features != n:
|
||||
setattr(model, name, nn.Linear(m.in_features, n))
|
||||
elif isinstance(m, nn.Sequential):
|
||||
types = [type(x) for x in m]
|
||||
if nn.Linear in types:
|
||||
i = types.index(nn.Linear) # nn.Linear index
|
||||
if m[i].out_features != n:
|
||||
m[i] = nn.Linear(m[i].in_features, n)
|
||||
elif nn.Conv2d in types:
|
||||
i = types.index(nn.Conv2d) # nn.Conv2d index
|
||||
if m[i].out_channels != n:
|
||||
m[i] = nn.Conv2d(m[i].in_channels, n, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def torch_distributed_zero_first(local_rank: int):
|
||||
# Decorator to make all processes in distributed training wait for each local_master to do something
|
||||
"""
|
||||
Decorator to make all processes in distributed training wait for each local_master to do something.
|
||||
"""
|
||||
if local_rank not in [-1, 0]:
|
||||
dist.barrier(device_ids=[local_rank])
|
||||
yield
|
||||
@@ -95,70 +38,69 @@ def torch_distributed_zero_first(local_rank: int):
|
||||
dist.barrier(device_ids=[0])
|
||||
|
||||
|
||||
def device_count():
|
||||
# Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Supports Linux and Windows
|
||||
assert platform.system() in ('Linux', 'Windows'), 'device_count() only supported on Linux or Windows'
|
||||
def date_modified(path=__file__):
|
||||
# return human-readable file modification date, i.e. '2021-3-26'
|
||||
t = datetime.datetime.fromtimestamp(Path(path).stat().st_mtime)
|
||||
return f'{t.year}-{t.month}-{t.day}'
|
||||
|
||||
|
||||
def git_describe(path=Path(__file__).parent): # path must be a directory
|
||||
# return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
|
||||
s = f'git -C {path} describe --tags --long --always'
|
||||
try:
|
||||
cmd = 'nvidia-smi -L | wc -l' if platform.system() == 'Linux' else 'nvidia-smi -L | find /c /v ""' # Windows
|
||||
return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])
|
||||
except Exception:
|
||||
return 0
|
||||
return subprocess.check_output(s, shell=True, stderr=subprocess.STDOUT).decode()[:-1]
|
||||
except subprocess.CalledProcessError as e:
|
||||
return '' # not a git repository
|
||||
|
||||
|
||||
def select_device(device='', batch_size=0, newline=True):
|
||||
# device = None or 'cpu' or 0 or '0' or '0,1,2,3'
|
||||
s = f'YOLOv3 🚀 {git_describe() or file_date()} Python-{platform.python_version()} torch-{torch.__version__} '
|
||||
device = str(device).strip().lower().replace('cuda:', '').replace('none', '') # to string, 'cuda:0' to '0'
|
||||
def select_device(device='', batch_size=None, newline=True):
|
||||
# device = 'cpu' or '0' or '0,1,2,3'
|
||||
s = f'YOLOv3 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
|
||||
device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0'
|
||||
cpu = device == 'cpu'
|
||||
mps = device == 'mps' # Apple Metal Performance Shaders (MPS)
|
||||
if cpu or mps:
|
||||
if cpu:
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
|
||||
elif device: # non-cpu device requested
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - must be before assert is_available()
|
||||
assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
|
||||
f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
|
||||
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
|
||||
|
||||
if not cpu and not mps and torch.cuda.is_available(): # prefer GPU if available
|
||||
cuda = not cpu and torch.cuda.is_available()
|
||||
if cuda:
|
||||
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
|
||||
n = len(devices) # device count
|
||||
if n > 1 and batch_size > 0: # check batch_size is divisible by device_count
|
||||
if n > 1 and batch_size: # check batch_size is divisible by device_count
|
||||
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
|
||||
space = ' ' * (len(s) + 1)
|
||||
for i, d in enumerate(devices):
|
||||
p = torch.cuda.get_device_properties(i)
|
||||
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n" # bytes to MB
|
||||
arg = 'cuda:0'
|
||||
elif mps and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available(): # prefer MPS if available
|
||||
s += 'MPS\n'
|
||||
arg = 'mps'
|
||||
else: # revert to CPU
|
||||
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2:.0f}MiB)\n" # bytes to MB
|
||||
else:
|
||||
s += 'CPU\n'
|
||||
arg = 'cpu'
|
||||
|
||||
if not newline:
|
||||
s = s.rstrip()
|
||||
LOGGER.info(s)
|
||||
return torch.device(arg)
|
||||
LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
|
||||
return torch.device('cuda:0' if cuda else 'cpu')
|
||||
|
||||
|
||||
def time_sync():
|
||||
# PyTorch-accurate time
|
||||
# pytorch-accurate time
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
return time.time()
|
||||
|
||||
|
||||
def profile(input, ops, n=10, device=None):
|
||||
""" YOLOv3 speed/memory/FLOPs profiler
|
||||
Usage:
|
||||
input = torch.randn(16, 3, 640, 640)
|
||||
m1 = lambda x: x * torch.sigmoid(x)
|
||||
m2 = nn.SiLU()
|
||||
profile(input, [m1, m2], n=100) # profile over 100 iterations
|
||||
"""
|
||||
# speed/memory/FLOPs profiler
|
||||
#
|
||||
# Usage:
|
||||
# input = torch.randn(16, 3, 640, 640)
|
||||
# m1 = lambda x: x * torch.sigmoid(x)
|
||||
# m2 = nn.SiLU()
|
||||
# profile(input, [m1, m2], n=100) # profile over 100 iterations
|
||||
|
||||
results = []
|
||||
if not isinstance(device, torch.device):
|
||||
device = select_device(device)
|
||||
device = device or select_device()
|
||||
print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
|
||||
f"{'input':>24s}{'output':>24s}")
|
||||
|
||||
@@ -171,7 +113,7 @@ def profile(input, ops, n=10, device=None):
|
||||
tf, tb, t = 0, 0, [0, 0, 0] # dt forward, backward
|
||||
try:
|
||||
flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPs
|
||||
except Exception:
|
||||
except:
|
||||
flops = 0
|
||||
|
||||
try:
|
||||
@@ -182,14 +124,15 @@ def profile(input, ops, n=10, device=None):
|
||||
try:
|
||||
_ = (sum(yi.sum() for yi in y) if isinstance(y, list) else y).sum().backward()
|
||||
t[2] = time_sync()
|
||||
except Exception: # no backward method
|
||||
except Exception as e: # no backward method
|
||||
# print(e) # for debug
|
||||
t[2] = float('nan')
|
||||
tf += (t[1] - t[0]) * 1000 / n # ms per op forward
|
||||
tb += (t[2] - t[1]) * 1000 / n # ms per op backward
|
||||
mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0 # (GB)
|
||||
s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' for x in (x, y)) # shapes
|
||||
p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0 # parameters
|
||||
s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
|
||||
s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
|
||||
p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
|
||||
print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}')
|
||||
results.append([p, flops, mem, tf, tb, s_in, s_out])
|
||||
except Exception as e:
|
||||
@@ -238,30 +181,30 @@ def sparsity(model):
|
||||
def prune(model, amount=0.3):
|
||||
# Prune model to requested global sparsity
|
||||
import torch.nn.utils.prune as prune
|
||||
print('Pruning model... ', end='')
|
||||
for name, m in model.named_modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
prune.l1_unstructured(m, name='weight', amount=amount) # prune
|
||||
prune.remove(m, 'weight') # make permanent
|
||||
LOGGER.info(f'Model pruned to {sparsity(model):.3g} global sparsity')
|
||||
print(' %.3g global sparsity' % sparsity(model))
|
||||
|
||||
|
||||
def fuse_conv_and_bn(conv, bn):
|
||||
# Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
|
||||
# Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
|
||||
fusedconv = nn.Conv2d(conv.in_channels,
|
||||
conv.out_channels,
|
||||
kernel_size=conv.kernel_size,
|
||||
stride=conv.stride,
|
||||
padding=conv.padding,
|
||||
dilation=conv.dilation,
|
||||
groups=conv.groups,
|
||||
bias=True).requires_grad_(False).to(conv.weight.device)
|
||||
|
||||
# Prepare filters
|
||||
# prepare filters
|
||||
w_conv = conv.weight.clone().view(conv.out_channels, -1)
|
||||
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
|
||||
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
|
||||
|
||||
# Prepare spatial bias
|
||||
# prepare spatial bias
|
||||
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
|
||||
b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
|
||||
fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
|
||||
@@ -269,7 +212,7 @@ def fuse_conv_and_bn(conv, bn):
|
||||
return fusedconv
|
||||
|
||||
|
||||
def model_info(model, verbose=False, imgsz=640):
|
||||
def model_info(model, verbose=False, img_size=640):
|
||||
# Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
|
||||
n_p = sum(x.numel() for x in model.parameters()) # number parameters
|
||||
n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
|
||||
@@ -281,29 +224,29 @@ def model_info(model, verbose=False, imgsz=640):
|
||||
(i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
|
||||
|
||||
try: # FLOPs
|
||||
p = next(model.parameters())
|
||||
stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32 # max stride
|
||||
im = torch.empty((1, p.shape[1], stride, stride), device=p.device) # input image in BCHW format
|
||||
flops = thop.profile(deepcopy(model), inputs=(im,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs
|
||||
imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz] # expand if int/float
|
||||
fs = f', {flops * imgsz[0] / stride * imgsz[1] / stride:.1f} GFLOPs' # 640x640 GFLOPs
|
||||
except Exception:
|
||||
from thop import profile
|
||||
stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
|
||||
img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
|
||||
flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs
|
||||
img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
|
||||
fs = ', %.1f GFLOPs' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPs
|
||||
except (ImportError, Exception):
|
||||
fs = ''
|
||||
|
||||
name = Path(model.yaml_file).stem.replace('yolov5', 'YOLOv3') if hasattr(model, 'yaml_file') else 'Model'
|
||||
LOGGER.info(f'{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}')
|
||||
LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
||||
|
||||
|
||||
def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|
||||
# Scales img(bs,3,y,x) by ratio constrained to gs-multiple
|
||||
# scales img(bs,3,y,x) by ratio constrained to gs-multiple
|
||||
if ratio == 1.0:
|
||||
return img
|
||||
h, w = img.shape[2:]
|
||||
s = (int(h * ratio), int(w * ratio)) # new size
|
||||
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
|
||||
if not same_shape: # pad/crop img
|
||||
h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
|
||||
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
|
||||
else:
|
||||
h, w = img.shape[2:]
|
||||
s = (int(h * ratio), int(w * ratio)) # new size
|
||||
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
|
||||
if not same_shape: # pad/crop img
|
||||
h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
|
||||
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
|
||||
|
||||
|
||||
def copy_attr(a, b, include=(), exclude=()):
|
||||
@@ -315,71 +258,8 @@ def copy_attr(a, b, include=(), exclude=()):
|
||||
setattr(a, k, v)
|
||||
|
||||
|
||||
def smart_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
|
||||
# YOLOv3 3-param group optimizer: 0) weights with decay, 1) weights no decay, 2) biases no decay
|
||||
g = [], [], [] # optimizer parameter groups
|
||||
bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k) # normalization layers, i.e. BatchNorm2d()
|
||||
for v in model.modules():
|
||||
for p_name, p in v.named_parameters(recurse=0):
|
||||
if p_name == 'bias': # bias (no decay)
|
||||
g[2].append(p)
|
||||
elif p_name == 'weight' and isinstance(v, bn): # weight (no decay)
|
||||
g[1].append(p)
|
||||
else:
|
||||
g[0].append(p) # weight (with decay)
|
||||
|
||||
if name == 'Adam':
|
||||
optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999)) # adjust beta1 to momentum
|
||||
elif name == 'AdamW':
|
||||
optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
|
||||
elif name == 'RMSProp':
|
||||
optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum)
|
||||
elif name == 'SGD':
|
||||
optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
|
||||
else:
|
||||
raise NotImplementedError(f'Optimizer {name} not implemented.')
|
||||
|
||||
optimizer.add_param_group({'params': g[0], 'weight_decay': decay}) # add g0 with weight_decay
|
||||
optimizer.add_param_group({'params': g[1], 'weight_decay': 0.0}) # add g1 (BatchNorm2d weights)
|
||||
LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "
|
||||
f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias')
|
||||
return optimizer
|
||||
|
||||
|
||||
def smart_hub_load(repo='ultralytics/yolov5', model='yolov5s', **kwargs):
|
||||
# YOLOv3 torch.hub.load() wrapper with smart error/issue handling
|
||||
if check_version(torch.__version__, '1.9.1'):
|
||||
kwargs['skip_validation'] = True # validation causes GitHub API rate limit errors
|
||||
if check_version(torch.__version__, '1.12.0'):
|
||||
kwargs['trust_repo'] = True # argument required starting in torch 0.12
|
||||
try:
|
||||
return torch.hub.load(repo, model, **kwargs)
|
||||
except Exception:
|
||||
return torch.hub.load(repo, model, force_reload=True, **kwargs)
|
||||
|
||||
|
||||
def smart_resume(ckpt, optimizer, ema=None, weights='yolov5s.pt', epochs=300, resume=True):
|
||||
# Resume training from a partially trained checkpoint
|
||||
best_fitness = 0.0
|
||||
start_epoch = ckpt['epoch'] + 1
|
||||
if ckpt['optimizer'] is not None:
|
||||
optimizer.load_state_dict(ckpt['optimizer']) # optimizer
|
||||
best_fitness = ckpt['best_fitness']
|
||||
if ema and ckpt.get('ema'):
|
||||
ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) # EMA
|
||||
ema.updates = ckpt['updates']
|
||||
if resume:
|
||||
assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.\n' \
|
||||
f"Start a new training without --resume, i.e. 'python train.py --weights {weights}'"
|
||||
LOGGER.info(f'Resuming training from {weights} from epoch {start_epoch} to {epochs} total epochs')
|
||||
if epochs < start_epoch:
|
||||
LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")
|
||||
epochs += ckpt['epoch'] # finetune additional epochs
|
||||
return best_fitness, start_epoch, epochs
|
||||
|
||||
|
||||
class EarlyStopping:
|
||||
# YOLOv3 simple early stopper
|
||||
# simple early stopper
|
||||
def __init__(self, patience=30):
|
||||
self.best_fitness = 0.0 # i.e. mAP
|
||||
self.best_epoch = 0
|
||||
@@ -402,30 +282,36 @@ class EarlyStopping:
|
||||
|
||||
|
||||
class ModelEMA:
|
||||
""" Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
|
||||
Keeps a moving average of everything in the model state_dict (parameters and buffers)
|
||||
For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
|
||||
""" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
|
||||
Keep a moving average of everything in the model state_dict (parameters and buffers).
|
||||
This is intended to allow functionality like
|
||||
https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
|
||||
A smoothed version of the weights is necessary for some training schemes to perform well.
|
||||
This class is sensitive where it is initialized in the sequence of model init,
|
||||
GPU assignment and distributed training wrappers.
|
||||
"""
|
||||
|
||||
def __init__(self, model, decay=0.9999, tau=2000, updates=0):
|
||||
def __init__(self, model, decay=0.9999, updates=0):
|
||||
# Create EMA
|
||||
self.ema = deepcopy(de_parallel(model)).eval() # FP32 EMA
|
||||
self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA
|
||||
# if next(model.parameters()).device.type != 'cpu':
|
||||
# self.ema.half() # FP16 EMA
|
||||
self.updates = updates # number of EMA updates
|
||||
self.decay = lambda x: decay * (1 - math.exp(-x / tau)) # decay exponential ramp (to help early epochs)
|
||||
self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
|
||||
for p in self.ema.parameters():
|
||||
p.requires_grad_(False)
|
||||
|
||||
def update(self, model):
|
||||
# Update EMA parameters
|
||||
self.updates += 1
|
||||
d = self.decay(self.updates)
|
||||
with torch.no_grad():
|
||||
self.updates += 1
|
||||
d = self.decay(self.updates)
|
||||
|
||||
msd = de_parallel(model).state_dict() # model state_dict
|
||||
for k, v in self.ema.state_dict().items():
|
||||
if v.dtype.is_floating_point: # true for FP16 and FP32
|
||||
v *= d
|
||||
v += (1 - d) * msd[k].detach()
|
||||
# assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype} and model {msd[k].dtype} must be FP32'
|
||||
msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict
|
||||
for k, v in self.ema.state_dict().items():
|
||||
if v.dtype.is_floating_point:
|
||||
v *= d
|
||||
v += (1 - d) * msd[k].detach()
|
||||
|
||||
def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
|
||||
# Update EMA attributes
|
||||
|
||||
Reference in New Issue
Block a user