341 lines
14 KiB
Python
341 lines
14 KiB
Python
|
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||
|
|
"""
|
||
|
|
PyTorch utils
|
||
|
|
"""
|
||
|
|
|
||
|
|
import datetime
|
||
|
|
import logging
|
||
|
|
import math
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
import subprocess
|
||
|
|
import time
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from copy import deepcopy
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.distributed as dist
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torchvision
|
||
|
|
|
||
|
|
try:
|
||
|
|
import thop # for FLOPs computation
|
||
|
|
except ImportError:
|
||
|
|
thop = None
|
||
|
|
|
||
|
|
LOGGER = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
@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.
|
||
|
|
"""
|
||
|
|
if local_rank not in [-1, 0]:
|
||
|
|
dist.barrier(device_ids=[local_rank])
|
||
|
|
yield
|
||
|
|
if local_rank == 0:
|
||
|
|
dist.barrier(device_ids=[0])
|
||
|
|
|
||
|
|
|
||
|
|
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:
|
||
|
|
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=None):
|
||
|
|
# device = 'cpu' or '0' or '0,1,2,3'
|
||
|
|
s = f'YOLOv5 🚀 {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'
|
||
|
|
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
|
||
|
|
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
|
||
|
|
|
||
|
|
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: # 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 / 1024 ** 2}MB)\n" # bytes to MB
|
||
|
|
else:
|
||
|
|
s += 'CPU\n'
|
||
|
|
|
||
|
|
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
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
return time.time()
|
||
|
|
|
||
|
|
|
||
|
|
def profile(input, ops, n=10, device=None):
|
||
|
|
# YOLOv5 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 = []
|
||
|
|
logging.basicConfig(format="%(message)s", level=logging.INFO)
|
||
|
|
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}")
|
||
|
|
|
||
|
|
for x in input if isinstance(input, list) else [input]:
|
||
|
|
x = x.to(device)
|
||
|
|
x.requires_grad = True
|
||
|
|
for m in ops if isinstance(ops, list) else [ops]:
|
||
|
|
m = m.to(device) if hasattr(m, 'to') else m # device
|
||
|
|
m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
|
||
|
|
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:
|
||
|
|
flops = 0
|
||
|
|
|
||
|
|
try:
|
||
|
|
for _ in range(n):
|
||
|
|
t[0] = time_sync()
|
||
|
|
y = m(x)
|
||
|
|
t[1] = time_sync()
|
||
|
|
try:
|
||
|
|
_ = (sum([yi.sum() for yi in y]) if isinstance(y, list) else y).sum().backward()
|
||
|
|
t[2] = time_sync()
|
||
|
|
except Exception as e: # no backward method
|
||
|
|
print(e)
|
||
|
|
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 = 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:
|
||
|
|
print(e)
|
||
|
|
results.append(None)
|
||
|
|
torch.cuda.empty_cache()
|
||
|
|
return results
|
||
|
|
|
||
|
|
|
||
|
|
def is_parallel(model):
|
||
|
|
# Returns True if model is of type DP or DDP
|
||
|
|
return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
|
||
|
|
|
||
|
|
|
||
|
|
def de_parallel(model):
|
||
|
|
# De-parallelize a model: returns single-GPU model if model is of type DP or DDP
|
||
|
|
return model.module if is_parallel(model) else model
|
||
|
|
|
||
|
|
|
||
|
|
def intersect_dicts(da, db, exclude=()):
|
||
|
|
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
|
||
|
|
return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}
|
||
|
|
|
||
|
|
|
||
|
|
def initialize_weights(model):
|
||
|
|
for m in model.modules():
|
||
|
|
t = type(m)
|
||
|
|
if t is nn.Conv2d:
|
||
|
|
pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||
|
|
elif t is nn.BatchNorm2d:
|
||
|
|
m.eps = 1e-3
|
||
|
|
m.momentum = 0.03
|
||
|
|
elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
|
||
|
|
m.inplace = True
|
||
|
|
|
||
|
|
|
||
|
|
def find_modules(model, mclass=nn.Conv2d):
|
||
|
|
# Finds layer indices matching module class 'mclass'
|
||
|
|
return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
|
||
|
|
|
||
|
|
|
||
|
|
def sparsity(model):
|
||
|
|
# Return global model sparsity
|
||
|
|
a, b = 0., 0.
|
||
|
|
for p in model.parameters():
|
||
|
|
a += p.numel()
|
||
|
|
b += (p == 0).sum()
|
||
|
|
return b / a
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
print(' %.3g global sparsity' % sparsity(model))
|
||
|
|
|
||
|
|
|
||
|
|
def fuse_conv_and_bn(conv, bn):
|
||
|
|
# 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,
|
||
|
|
groups=conv.groups,
|
||
|
|
bias=True).requires_grad_(False).to(conv.weight.device)
|
||
|
|
|
||
|
|
# 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
|
||
|
|
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)
|
||
|
|
|
||
|
|
return fusedconv
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
if verbose:
|
||
|
|
print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
|
||
|
|
for i, (name, p) in enumerate(model.named_parameters()):
|
||
|
|
name = name.replace('module_list.', '')
|
||
|
|
print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
|
||
|
|
(i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
|
||
|
|
|
||
|
|
try: # FLOPs
|
||
|
|
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 = ''
|
||
|
|
|
||
|
|
LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
||
|
|
|
||
|
|
|
||
|
|
def load_classifier(name='resnet101', n=2):
|
||
|
|
# Loads a pretrained model reshaped to n-class output
|
||
|
|
model = torchvision.models.__dict__[name](pretrained=True)
|
||
|
|
|
||
|
|
# ResNet model properties
|
||
|
|
# input_size = [3, 224, 224]
|
||
|
|
# input_space = 'RGB'
|
||
|
|
# input_range = [0, 1]
|
||
|
|
# mean = [0.485, 0.456, 0.406]
|
||
|
|
# std = [0.229, 0.224, 0.225]
|
||
|
|
|
||
|
|
# Reshape output to n classes
|
||
|
|
filters = model.fc.weight.shape[1]
|
||
|
|
model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
|
||
|
|
model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
|
||
|
|
model.fc.out_features = n
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
if ratio == 1.0:
|
||
|
|
return img
|
||
|
|
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=()):
|
||
|
|
# Copy attributes from b to a, options to only include [...] and to exclude [...]
|
||
|
|
for k, v in b.__dict__.items():
|
||
|
|
if (len(include) and k not in include) or k.startswith('_') or k in exclude:
|
||
|
|
continue
|
||
|
|
else:
|
||
|
|
setattr(a, k, v)
|
||
|
|
|
||
|
|
|
||
|
|
class EarlyStopping:
|
||
|
|
# YOLOv5 simple early stopper
|
||
|
|
def __init__(self, patience=30):
|
||
|
|
self.best_fitness = 0.0 # i.e. mAP
|
||
|
|
self.best_epoch = 0
|
||
|
|
self.patience = patience or float('inf') # epochs to wait after fitness stops improving to stop
|
||
|
|
self.possible_stop = False # possible stop may occur next epoch
|
||
|
|
|
||
|
|
def __call__(self, epoch, fitness):
|
||
|
|
if fitness >= self.best_fitness: # >= 0 to allow for early zero-fitness stage of training
|
||
|
|
self.best_epoch = epoch
|
||
|
|
self.best_fitness = fitness
|
||
|
|
delta = epoch - self.best_epoch # epochs without improvement
|
||
|
|
self.possible_stop = delta >= (self.patience - 1) # possible stop may occur next epoch
|
||
|
|
stop = delta >= self.patience # stop training if patience exceeded
|
||
|
|
if stop:
|
||
|
|
LOGGER.info(f'EarlyStopping patience {self.patience} exceeded, stopping training.')
|
||
|
|
return stop
|
||
|
|
|
||
|
|
|
||
|
|
class ModelEMA:
|
||
|
|
""" 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, updates=0):
|
||
|
|
# Create 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 / 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
|
||
|
|
with torch.no_grad():
|
||
|
|
self.updates += 1
|
||
|
|
d = self.decay(self.updates)
|
||
|
|
|
||
|
|
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
|
||
|
|
copy_attr(self.ema, model, include, exclude)
|