您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

341 行
14KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. PyTorch utils
  4. """
  5. import datetime
  6. import logging
  7. import math
  8. import os
  9. import platform
  10. import subprocess
  11. import time
  12. from contextlib import contextmanager
  13. from copy import deepcopy
  14. from pathlib import Path
  15. import torch
  16. import torch.distributed as dist
  17. import torch.nn as nn
  18. import torch.nn.functional as F
  19. import torchvision
  20. try:
  21. import thop # for FLOPs computation
  22. except ImportError:
  23. thop = None
  24. LOGGER = logging.getLogger(__name__)
  25. @contextmanager
  26. def torch_distributed_zero_first(local_rank: int):
  27. """
  28. Decorator to make all processes in distributed training wait for each local_master to do something.
  29. """
  30. if local_rank not in [-1, 0]:
  31. dist.barrier(device_ids=[local_rank])
  32. yield
  33. if local_rank == 0:
  34. dist.barrier(device_ids=[0])
  35. def date_modified(path=__file__):
  36. # return human-readable file modification date, i.e. '2021-3-26'
  37. t = datetime.datetime.fromtimestamp(Path(path).stat().st_mtime)
  38. return f'{t.year}-{t.month}-{t.day}'
  39. def git_describe(path=Path(__file__).parent): # path must be a directory
  40. # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
  41. s = f'git -C {path} describe --tags --long --always'
  42. try:
  43. return subprocess.check_output(s, shell=True, stderr=subprocess.STDOUT).decode()[:-1]
  44. except subprocess.CalledProcessError as e:
  45. return '' # not a git repository
  46. def select_device(device='', batch_size=None):
  47. # device = 'cpu' or '0' or '0,1,2,3'
  48. s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
  49. device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0'
  50. cpu = device == 'cpu'
  51. if cpu:
  52. os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
  53. elif device: # non-cpu device requested
  54. os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
  55. assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
  56. cuda = not cpu and torch.cuda.is_available()
  57. if cuda:
  58. devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
  59. n = len(devices) # device count
  60. if n > 1 and batch_size: # check batch_size is divisible by device_count
  61. assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
  62. space = ' ' * (len(s) + 1)
  63. for i, d in enumerate(devices):
  64. p = torch.cuda.get_device_properties(i)
  65. s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
  66. else:
  67. s += 'CPU\n'
  68. LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
  69. return torch.device('cuda:0' if cuda else 'cpu')
  70. def time_sync():
  71. # pytorch-accurate time
  72. if torch.cuda.is_available():
  73. torch.cuda.synchronize()
  74. return time.time()
  75. def profile(input, ops, n=10, device=None):
  76. # YOLOv5 speed/memory/FLOPs profiler
  77. #
  78. # Usage:
  79. # input = torch.randn(16, 3, 640, 640)
  80. # m1 = lambda x: x * torch.sigmoid(x)
  81. # m2 = nn.SiLU()
  82. # profile(input, [m1, m2], n=100) # profile over 100 iterations
  83. results = []
  84. logging.basicConfig(format="%(message)s", level=logging.INFO)
  85. device = device or select_device()
  86. print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
  87. f"{'input':>24s}{'output':>24s}")
  88. for x in input if isinstance(input, list) else [input]:
  89. x = x.to(device)
  90. x.requires_grad = True
  91. for m in ops if isinstance(ops, list) else [ops]:
  92. m = m.to(device) if hasattr(m, 'to') else m # device
  93. m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
  94. tf, tb, t = 0., 0., [0., 0., 0.] # dt forward, backward
  95. try:
  96. flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPs
  97. except:
  98. flops = 0
  99. try:
  100. for _ in range(n):
  101. t[0] = time_sync()
  102. y = m(x)
  103. t[1] = time_sync()
  104. try:
  105. _ = (sum([yi.sum() for yi in y]) if isinstance(y, list) else y).sum().backward()
  106. t[2] = time_sync()
  107. except Exception as e: # no backward method
  108. print(e)
  109. t[2] = float('nan')
  110. tf += (t[1] - t[0]) * 1000 / n # ms per op forward
  111. tb += (t[2] - t[1]) * 1000 / n # ms per op backward
  112. mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0 # (GB)
  113. s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
  114. s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
  115. p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
  116. print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}')
  117. results.append([p, flops, mem, tf, tb, s_in, s_out])
  118. except Exception as e:
  119. print(e)
  120. results.append(None)
  121. torch.cuda.empty_cache()
  122. return results
  123. def is_parallel(model):
  124. # Returns True if model is of type DP or DDP
  125. return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
  126. def de_parallel(model):
  127. # De-parallelize a model: returns single-GPU model if model is of type DP or DDP
  128. return model.module if is_parallel(model) else model
  129. def intersect_dicts(da, db, exclude=()):
  130. # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
  131. 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}
  132. def initialize_weights(model):
  133. for m in model.modules():
  134. t = type(m)
  135. if t is nn.Conv2d:
  136. pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
  137. elif t is nn.BatchNorm2d:
  138. m.eps = 1e-3
  139. m.momentum = 0.03
  140. elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
  141. m.inplace = True
  142. def find_modules(model, mclass=nn.Conv2d):
  143. # Finds layer indices matching module class 'mclass'
  144. return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
  145. def sparsity(model):
  146. # Return global model sparsity
  147. a, b = 0., 0.
  148. for p in model.parameters():
  149. a += p.numel()
  150. b += (p == 0).sum()
  151. return b / a
  152. def prune(model, amount=0.3):
  153. # Prune model to requested global sparsity
  154. import torch.nn.utils.prune as prune
  155. print('Pruning model... ', end='')
  156. for name, m in model.named_modules():
  157. if isinstance(m, nn.Conv2d):
  158. prune.l1_unstructured(m, name='weight', amount=amount) # prune
  159. prune.remove(m, 'weight') # make permanent
  160. print(' %.3g global sparsity' % sparsity(model))
  161. def fuse_conv_and_bn(conv, bn):
  162. # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
  163. fusedconv = nn.Conv2d(conv.in_channels,
  164. conv.out_channels,
  165. kernel_size=conv.kernel_size,
  166. stride=conv.stride,
  167. padding=conv.padding,
  168. groups=conv.groups,
  169. bias=True).requires_grad_(False).to(conv.weight.device)
  170. # prepare filters
  171. w_conv = conv.weight.clone().view(conv.out_channels, -1)
  172. w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
  173. fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
  174. # prepare spatial bias
  175. b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
  176. b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
  177. fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
  178. return fusedconv
  179. def model_info(model, verbose=False, img_size=640):
  180. # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
  181. n_p = sum(x.numel() for x in model.parameters()) # number parameters
  182. n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
  183. if verbose:
  184. print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
  185. for i, (name, p) in enumerate(model.named_parameters()):
  186. name = name.replace('module_list.', '')
  187. print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
  188. (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
  189. try: # FLOPs
  190. from thop import profile
  191. stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
  192. img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
  193. flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs
  194. img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
  195. fs = ', %.1f GFLOPs' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPs
  196. except (ImportError, Exception):
  197. fs = ''
  198. LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
  199. def load_classifier(name='resnet101', n=2):
  200. # Loads a pretrained model reshaped to n-class output
  201. model = torchvision.models.__dict__[name](pretrained=True)
  202. # ResNet model properties
  203. # input_size = [3, 224, 224]
  204. # input_space = 'RGB'
  205. # input_range = [0, 1]
  206. # mean = [0.485, 0.456, 0.406]
  207. # std = [0.229, 0.224, 0.225]
  208. # Reshape output to n classes
  209. filters = model.fc.weight.shape[1]
  210. model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
  211. model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
  212. model.fc.out_features = n
  213. return model
  214. def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
  215. # scales img(bs,3,y,x) by ratio constrained to gs-multiple
  216. if ratio == 1.0:
  217. return img
  218. else:
  219. h, w = img.shape[2:]
  220. s = (int(h * ratio), int(w * ratio)) # new size
  221. img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
  222. if not same_shape: # pad/crop img
  223. h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
  224. return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
  225. def copy_attr(a, b, include=(), exclude=()):
  226. # Copy attributes from b to a, options to only include [...] and to exclude [...]
  227. for k, v in b.__dict__.items():
  228. if (len(include) and k not in include) or k.startswith('_') or k in exclude:
  229. continue
  230. else:
  231. setattr(a, k, v)
  232. class EarlyStopping:
  233. # YOLOv5 simple early stopper
  234. def __init__(self, patience=30):
  235. self.best_fitness = 0.0 # i.e. mAP
  236. self.best_epoch = 0
  237. self.patience = patience or float('inf') # epochs to wait after fitness stops improving to stop
  238. self.possible_stop = False # possible stop may occur next epoch
  239. def __call__(self, epoch, fitness):
  240. if fitness >= self.best_fitness: # >= 0 to allow for early zero-fitness stage of training
  241. self.best_epoch = epoch
  242. self.best_fitness = fitness
  243. delta = epoch - self.best_epoch # epochs without improvement
  244. self.possible_stop = delta >= (self.patience - 1) # possible stop may occur next epoch
  245. stop = delta >= self.patience # stop training if patience exceeded
  246. if stop:
  247. LOGGER.info(f'EarlyStopping patience {self.patience} exceeded, stopping training.')
  248. return stop
  249. class ModelEMA:
  250. """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
  251. Keep a moving average of everything in the model state_dict (parameters and buffers).
  252. This is intended to allow functionality like
  253. https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
  254. A smoothed version of the weights is necessary for some training schemes to perform well.
  255. This class is sensitive where it is initialized in the sequence of model init,
  256. GPU assignment and distributed training wrappers.
  257. """
  258. def __init__(self, model, decay=0.9999, updates=0):
  259. # Create EMA
  260. self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA
  261. # if next(model.parameters()).device.type != 'cpu':
  262. # self.ema.half() # FP16 EMA
  263. self.updates = updates # number of EMA updates
  264. self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
  265. for p in self.ema.parameters():
  266. p.requires_grad_(False)
  267. def update(self, model):
  268. # Update EMA parameters
  269. with torch.no_grad():
  270. self.updates += 1
  271. d = self.decay(self.updates)
  272. msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict
  273. for k, v in self.ema.state_dict().items():
  274. if v.dtype.is_floating_point:
  275. v *= d
  276. v += (1. - d) * msd[k].detach()
  277. def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
  278. # Update EMA attributes
  279. copy_attr(self.ema, model, include, exclude)