car-detection-bayes/utils/torch_utils.py

195 lines
8.1 KiB
Python
Raw Normal View History

2020-03-29 20:14:54 +00:00
import math
2019-09-10 12:59:45 +00:00
import os
2020-03-04 18:26:35 +00:00
import time
2020-03-14 03:12:54 +00:00
from copy import deepcopy
2019-09-13 14:00:52 +00:00
import torch
2020-02-28 06:50:26 +00:00
import torch.backends.cudnn as cudnn
2020-03-14 03:12:54 +00:00
import torch.nn as nn
2020-03-16 01:39:54 +00:00
import torch.nn.functional as F
def init_seeds(seed=0):
torch.manual_seed(seed)
2019-09-10 08:56:56 +00:00
# Remove randomness (may be slower on Tesla GPUs) # https://pytorch.org/docs/stable/notes/randomness.html
if seed == 0:
2020-02-28 06:50:26 +00:00
cudnn.deterministic = True
cudnn.benchmark = False
2019-11-25 04:29:29 +00:00
def select_device(device='', apex=False, batch_size=None):
2019-09-26 11:52:37 +00:00
# device = 'cpu' or '0' or '0,1,2,3'
cpu_request = device.lower() == 'cpu'
if device and not cpu_request: # if device requested other than 'cpu'
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity
cuda = False if cpu_request else torch.cuda.is_available()
2019-04-08 13:41:14 +00:00
if cuda:
c = 1024 ** 2 # bytes to MB
ng = torch.cuda.device_count()
2019-11-25 04:29:29 +00:00
if ng > 1 and batch_size: # check that batch_size is compatible with device_count
assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)
2019-04-08 13:41:14 +00:00
x = [torch.cuda.get_device_properties(i) for i in range(ng)]
2019-11-25 04:29:29 +00:00
s = 'Using CUDA ' + ('Apex ' if apex else '') # apex for mixed precision https://github.com/NVIDIA/apex
2019-07-16 17:09:40 +00:00
for i in range(0, ng):
2019-07-16 17:10:33 +00:00
if i == 1:
2019-11-25 04:29:29 +00:00
s = ' ' * len(s)
2019-07-16 17:09:40 +00:00
print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
2019-11-25 04:31:06 +00:00
(s, i, x[i].name, x[i].total_memory / c))
2019-09-26 11:52:37 +00:00
else:
print('Using CPU')
2019-02-16 13:33:52 +00:00
2019-05-03 16:14:16 +00:00
print('') # skip a line
2019-09-26 11:52:37 +00:00
return torch.device('cuda:0' if cuda else 'cpu')
2019-04-19 18:41:18 +00:00
2020-03-04 18:26:35 +00:00
def time_synchronized():
torch.cuda.synchronize() if torch.cuda.is_available() else None
return time.time()
2020-04-11 17:45:33 +00:00
def initialize_weights(model):
for m in model.modules():
2020-04-14 08:20:57 +00:00
t = type(m)
if t is nn.Conv2d:
2020-04-14 22:58:32 +00:00
pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
2020-04-14 08:20:57 +00:00
elif t is nn.BatchNorm2d:
2020-04-13 01:44:18 +00:00
m.eps = 1e-4
m.momentum = 0.03
2020-04-14 08:20:57 +00:00
elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
m.inplace = True
2020-04-11 17:45:33 +00:00
2020-04-14 00:48:30 +00:00
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)]
2019-04-19 18:41:18 +00:00
def fuse_conv_and_bn(conv, bn):
# https://tehnokv.com/posts/fusing-batchnorm-and-conv/
with torch.no_grad():
# init
2019-07-24 17:02:24 +00:00
fusedconv = torch.nn.Conv2d(conv.in_channels,
conv.out_channels,
kernel_size=conv.kernel_size,
stride=conv.stride,
padding=conv.padding,
bias=True)
2019-04-19 18:41:18 +00:00
# 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.size()))
# prepare spatial bias
if conv.bias is not None:
b_conv = conv.bias
else:
b_conv = torch.zeros(conv.weight.size(0))
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)
2019-04-19 18:41:18 +00:00
return fusedconv
2019-10-06 22:50:47 +00:00
2020-03-14 23:46:54 +00:00
def model_info(model, verbose=False):
2019-10-06 22:50:47 +00:00
# Plots a line-by-line description of a PyTorch model
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
2020-03-14 23:46:54 +00:00
if verbose:
2019-10-06 22:50:47 +00:00
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()))
2019-10-10 12:40:18 +00:00
2020-04-01 21:05:41 +00:00
try: # FLOPS
from thop import profile
2020-04-13 01:44:18 +00:00
macs, _ = profile(model, inputs=(torch.zeros(1, 3, 480, 640),))
2020-04-01 21:05:41 +00:00
fs = ', %.1f GFLOPS' % (macs / 1E9 * 2)
except:
fs = ''
print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs))
2020-03-19 19:30:07 +00:00
2019-10-10 12:40:18 +00:00
def load_classifier(name='resnet101', n=2):
# Loads a pretrained model reshaped to n-class output
import pretrainedmodels # https://github.com/Cadene/pretrained-models.pytorch#torchvision
model = pretrainedmodels.__dict__[name](num_classes=1000, pretrained='imagenet')
# Display model properties
for x in ['model.input_size', 'model.input_space', 'model.input_range', 'model.mean', 'model.std']:
print(x + ' =', eval(x))
# Reshape output to n classes
filters = model.last_linear.weight.shape[1]
model.last_linear.bias = torch.nn.Parameter(torch.zeros(n))
model.last_linear.weight = torch.nn.Parameter(torch.zeros(n, filters))
model.last_linear.out_features = n
return model
2020-03-14 03:12:54 +00:00
2020-04-07 19:51:52 +00:00
def scale_img(img, ratio=1.0, same_shape=True): # img(16,3,256,416), r=ratio
# scales img(bs,3,y,x) by ratio
2020-03-16 01:39:54 +00:00
h, w = img.shape[2:]
2020-04-07 19:51:52 +00:00
s = (int(h * ratio), int(w * ratio)) # new size
2020-03-16 01:39:54 +00:00
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
2020-04-07 19:51:52 +00:00
if not same_shape: # pad/crop img
2020-04-07 20:35:47 +00:00
gs = 64 # (pixels) grid size
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
2020-03-16 01:39:54 +00:00
2020-03-14 03:12:54 +00:00
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.
E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use
RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA
smoothing of weights to match results. Pay attention to the decay constant you are using
relative to your update count per epoch.
To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but
disable validation of the EMA weights. Validation will have to be done manually in a separate
process, or after the training stops converging.
This class is sensitive where it is initialized in the sequence of model init,
GPU assignment and distributed training wrappers.
I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and single-GPU.
"""
2020-03-29 20:14:54 +00:00
def __init__(self, model, decay=0.9999, device=''):
2020-03-14 03:12:54 +00:00
# make a copy of the model for accumulating moving average of weights
self.ema = deepcopy(model)
self.ema.eval()
2020-03-29 20:14:54 +00:00
self.updates = 0 # number of EMA updates
2020-04-09 04:01:58 +00:00
self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
2020-03-14 03:12:54 +00:00
self.device = device # perform ema on different device from model if set
if device:
self.ema.to(device=device)
for p in self.ema.parameters():
p.requires_grad_(False)
def update(self, model):
2020-03-29 20:14:54 +00:00
self.updates += 1
d = self.decay(self.updates)
2020-03-14 03:12:54 +00:00
with torch.no_grad():
if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel):
msd, esd = model.module.state_dict(), self.ema.module.state_dict()
else:
msd, esd = model.state_dict(), self.ema.state_dict()
2020-03-17 00:51:40 +00:00
2020-03-14 23:23:14 +00:00
for k, v in esd.items():
if v.dtype.is_floating_point:
v *= d
v += (1. - d) * msd[k].detach()
def update_attr(self, model):
# Assign attributes (which may change during training)
for k in model.__dict__.keys():
if not k.startswith('_'):
2020-03-17 00:51:40 +00:00
setattr(self.ema, k, getattr(model, k))