torch_utils.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374
  1. # YOLOR PyTorch utils
  2. import datetime
  3. import logging
  4. import math
  5. import os
  6. import platform
  7. import subprocess
  8. import time
  9. from contextlib import contextmanager
  10. from copy import deepcopy
  11. from pathlib import Path
  12. import torch
  13. import torch.backends.cudnn as cudnn
  14. import torch.nn as nn
  15. import torch.nn.functional as F
  16. import torchvision
  17. try:
  18. import thop # for FLOPS computation
  19. except ImportError:
  20. thop = None
  21. logger = logging.getLogger(__name__)
  22. @contextmanager
  23. def torch_distributed_zero_first(local_rank: int):
  24. """
  25. Decorator to make all processes in distributed training wait for each local_master to do something.
  26. """
  27. if local_rank not in [-1, 0]:
  28. torch.distributed.barrier()
  29. yield
  30. if local_rank == 0:
  31. torch.distributed.barrier()
  32. def init_torch_seeds(seed=0):
  33. # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
  34. torch.manual_seed(seed)
  35. if seed == 0: # slower, more reproducible
  36. cudnn.benchmark, cudnn.deterministic = False, True
  37. else: # faster, less reproducible
  38. cudnn.benchmark, cudnn.deterministic = True, False
  39. def date_modified(path=__file__):
  40. # return human-readable file modification date, i.e. '2021-3-26'
  41. t = datetime.datetime.fromtimestamp(Path(path).stat().st_mtime)
  42. return f'{t.year}-{t.month}-{t.day}'
  43. def git_describe(path=Path(__file__).parent): # path must be a directory
  44. # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
  45. s = f'git -C {path} describe --tags --long --always'
  46. try:
  47. return subprocess.check_output(s, shell=True, stderr=subprocess.STDOUT).decode()[:-1]
  48. except subprocess.CalledProcessError as e:
  49. return '' # not a git repository
  50. def select_device(device='', batch_size=None):
  51. # device = 'cpu' or '0' or '0,1,2,3'
  52. s = f'YOLOR 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
  53. cpu = device.lower() == 'cpu'
  54. if cpu:
  55. os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
  56. elif device: # non-cpu device requested
  57. # os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
  58. assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
  59. cuda = not cpu and torch.cuda.is_available()
  60. if cuda:
  61. n = torch.cuda.device_count()
  62. if n > 1 and batch_size: # check that batch_size is compatible with device_count
  63. assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
  64. space = ' ' * len(s)
  65. for i, d in enumerate(device.split(',') if device else range(n)):
  66. p = torch.cuda.get_device_properties(i)
  67. s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
  68. else:
  69. s += 'CPU\n'
  70. logger.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
  71. return torch.device('cuda:0' if cuda else 'cpu')
  72. def time_synchronized():
  73. # pytorch-accurate time
  74. if torch.cuda.is_available():
  75. torch.cuda.synchronize()
  76. return time.time()
  77. def profile(x, ops, n=100, device=None):
  78. # profile a pytorch module or list of modules. Example usage:
  79. # x = torch.randn(16, 3, 640, 640) # input
  80. # m1 = lambda x: x * torch.sigmoid(x)
  81. # m2 = nn.SiLU()
  82. # profile(x, [m1, m2], n=100) # profile speed over 100 iterations
  83. device = device or torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  84. x = x.to(device)
  85. x.requires_grad = True
  86. print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '')
  87. print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}")
  88. for m in ops if isinstance(ops, list) else [ops]:
  89. m = m.to(device) if hasattr(m, 'to') else m # device
  90. m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m # type
  91. dtf, dtb, t = 0., 0., [0., 0., 0.] # dt forward, backward
  92. try:
  93. flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPS
  94. except:
  95. flops = 0
  96. for _ in range(n):
  97. t[0] = time_synchronized()
  98. y = m(x)
  99. t[1] = time_synchronized()
  100. try:
  101. _ = y.sum().backward()
  102. t[2] = time_synchronized()
  103. except: # no backward method
  104. t[2] = float('nan')
  105. dtf += (t[1] - t[0]) * 1000 / n # ms per op forward
  106. dtb += (t[2] - t[1]) * 1000 / n # ms per op backward
  107. s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
  108. s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
  109. p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
  110. print(f'{p:12}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
  111. def is_parallel(model):
  112. return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
  113. def intersect_dicts(da, db, exclude=()):
  114. # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
  115. 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}
  116. def initialize_weights(model):
  117. for m in model.modules():
  118. t = type(m)
  119. if t is nn.Conv2d:
  120. pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
  121. elif t is nn.BatchNorm2d:
  122. m.eps = 1e-3
  123. m.momentum = 0.03
  124. elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
  125. m.inplace = True
  126. def find_modules(model, mclass=nn.Conv2d):
  127. # Finds layer indices matching module class 'mclass'
  128. return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
  129. def sparsity(model):
  130. # Return global model sparsity
  131. a, b = 0., 0.
  132. for p in model.parameters():
  133. a += p.numel()
  134. b += (p == 0).sum()
  135. return b / a
  136. def prune(model, amount=0.3):
  137. # Prune model to requested global sparsity
  138. import torch.nn.utils.prune as prune
  139. print('Pruning model... ', end='')
  140. for name, m in model.named_modules():
  141. if isinstance(m, nn.Conv2d):
  142. prune.l1_unstructured(m, name='weight', amount=amount) # prune
  143. prune.remove(m, 'weight') # make permanent
  144. print(' %.3g global sparsity' % sparsity(model))
  145. def fuse_conv_and_bn(conv, bn):
  146. # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
  147. fusedconv = nn.Conv2d(conv.in_channels,
  148. conv.out_channels,
  149. kernel_size=conv.kernel_size,
  150. stride=conv.stride,
  151. padding=conv.padding,
  152. groups=conv.groups,
  153. bias=True).requires_grad_(False).to(conv.weight.device)
  154. # prepare filters
  155. w_conv = conv.weight.clone().view(conv.out_channels, -1)
  156. w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
  157. fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
  158. # prepare spatial bias
  159. b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
  160. b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
  161. fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
  162. return fusedconv
  163. def model_info(model, verbose=False, img_size=640):
  164. # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
  165. n_p = sum(x.numel() for x in model.parameters()) # number parameters
  166. n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
  167. if verbose:
  168. print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
  169. for i, (name, p) in enumerate(model.named_parameters()):
  170. name = name.replace('module_list.', '')
  171. print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
  172. (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
  173. try: # FLOPS
  174. from thop import profile
  175. stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
  176. img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
  177. flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS
  178. img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
  179. fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPS
  180. except (ImportError, Exception):
  181. fs = ''
  182. logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
  183. def load_classifier(name='resnet101', n=2):
  184. # Loads a pretrained model reshaped to n-class output
  185. model = torchvision.models.__dict__[name](pretrained=True)
  186. # ResNet model properties
  187. # input_size = [3, 224, 224]
  188. # input_space = 'RGB'
  189. # input_range = [0, 1]
  190. # mean = [0.485, 0.456, 0.406]
  191. # std = [0.229, 0.224, 0.225]
  192. # Reshape output to n classes
  193. filters = model.fc.weight.shape[1]
  194. model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
  195. model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
  196. model.fc.out_features = n
  197. return model
  198. def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
  199. # scales img(bs,3,y,x) by ratio constrained to gs-multiple
  200. if ratio == 1.0:
  201. return img
  202. else:
  203. h, w = img.shape[2:]
  204. s = (int(h * ratio), int(w * ratio)) # new size
  205. img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
  206. if not same_shape: # pad/crop img
  207. h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
  208. return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
  209. def copy_attr(a, b, include=(), exclude=()):
  210. # Copy attributes from b to a, options to only include [...] and to exclude [...]
  211. for k, v in b.__dict__.items():
  212. if (len(include) and k not in include) or k.startswith('_') or k in exclude:
  213. continue
  214. else:
  215. setattr(a, k, v)
  216. class ModelEMA:
  217. """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
  218. Keep a moving average of everything in the model state_dict (parameters and buffers).
  219. This is intended to allow functionality like
  220. https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
  221. A smoothed version of the weights is necessary for some training schemes to perform well.
  222. This class is sensitive where it is initialized in the sequence of model init,
  223. GPU assignment and distributed training wrappers.
  224. """
  225. def __init__(self, model, decay=0.9999, updates=0):
  226. # Create EMA
  227. self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA
  228. # if next(model.parameters()).device.type != 'cpu':
  229. # self.ema.half() # FP16 EMA
  230. self.updates = updates # number of EMA updates
  231. self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
  232. for p in self.ema.parameters():
  233. p.requires_grad_(False)
  234. def update(self, model):
  235. # Update EMA parameters
  236. with torch.no_grad():
  237. self.updates += 1
  238. d = self.decay(self.updates)
  239. msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict
  240. for k, v in self.ema.state_dict().items():
  241. if v.dtype.is_floating_point:
  242. v *= d
  243. v += (1. - d) * msd[k].detach()
  244. def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
  245. # Update EMA attributes
  246. copy_attr(self.ema, model, include, exclude)
  247. class BatchNormXd(torch.nn.modules.batchnorm._BatchNorm):
  248. def _check_input_dim(self, input):
  249. # The only difference between BatchNorm1d, BatchNorm2d, BatchNorm3d, etc
  250. # is this method that is overwritten by the sub-class
  251. # This original goal of this method was for tensor sanity checks
  252. # If you're ok bypassing those sanity checks (eg. if you trust your inference
  253. # to provide the right dimensional inputs), then you can just use this method
  254. # for easy conversion from SyncBatchNorm
  255. # (unfortunately, SyncBatchNorm does not store the original class - if it did
  256. # we could return the one that was originally created)
  257. return
  258. def revert_sync_batchnorm(module):
  259. # this is very similar to the function that it is trying to revert:
  260. # https://github.com/pytorch/pytorch/blob/c8b3686a3e4ba63dc59e5dcfe5db3430df256833/torch/nn/modules/batchnorm.py#L679
  261. module_output = module
  262. if isinstance(module, torch.nn.modules.batchnorm.SyncBatchNorm):
  263. new_cls = BatchNormXd
  264. module_output = BatchNormXd(module.num_features,
  265. module.eps, module.momentum,
  266. module.affine,
  267. module.track_running_stats)
  268. if module.affine:
  269. with torch.no_grad():
  270. module_output.weight = module.weight
  271. module_output.bias = module.bias
  272. module_output.running_mean = module.running_mean
  273. module_output.running_var = module.running_var
  274. module_output.num_batches_tracked = module.num_batches_tracked
  275. if hasattr(module, "qconfig"):
  276. module_output.qconfig = module.qconfig
  277. for name, child in module.named_children():
  278. module_output.add_module(name, revert_sync_batchnorm(child))
  279. del module
  280. return module_output
  281. class TracedModel(nn.Module):
  282. def __init__(self, model=None, device=None, img_size=(640,640)):
  283. super(TracedModel, self).__init__()
  284. print(" Convert model to Traced-model... ")
  285. self.stride = model.stride
  286. self.names = model.names
  287. self.model = model
  288. self.model = revert_sync_batchnorm(self.model)
  289. self.model.to('cpu')
  290. self.model.eval()
  291. self.detect_layer = self.model.model[-1]
  292. self.model.traced = True
  293. rand_example = torch.rand(1, 3, img_size, img_size)
  294. traced_script_module = torch.jit.trace(self.model, rand_example, strict=False)
  295. #traced_script_module = torch.jit.script(self.model)
  296. traced_script_module.save("traced_model.pt")
  297. print(" traced_script_module saved! ")
  298. self.model = traced_script_module
  299. self.model.to(device)
  300. self.detect_layer.to(device)
  301. print(" model is traced! \n")
  302. def forward(self, x, augment=False, profile=False):
  303. out = self.model(x)
  304. out = self.detect_layer(out)
  305. return out