|
- import os
-
- import torch
-
-
- def init_seeds(seed=0):
- torch.manual_seed(seed)
- torch.cuda.manual_seed(seed)
- torch.cuda.manual_seed_all(seed)
-
- # Remove randomness (may be slower on Tesla GPUs) # https://pytorch.org/docs/stable/notes/randomness.html
- if seed == 0:
- torch.backends.cudnn.deterministic = True
- torch.backends.cudnn.benchmark = False
-
-
- def select_device(device='', apex=False):
- # 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()
- if cuda:
- c = 1024 ** 2 # bytes to MB
- ng = torch.cuda.device_count()
- x = [torch.cuda.get_device_properties(i) for i in range(ng)]
- cuda_str = 'Using CUDA ' + ('Apex ' if apex else '') # apex for mixed precision https://github.com/NVIDIA/apex
- for i in range(0, ng):
- if i == 1:
- cuda_str = ' ' * len(cuda_str)
- print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
- (cuda_str, i, x[i].name, x[i].total_memory / c))
- else:
- print('Using CPU')
-
- print('') # skip a line
- return torch.device('cuda:0' if cuda else 'cpu')
-
-
- def fuse_conv_and_bn(conv, bn):
- # https://tehnokv.com/posts/fusing-batchnorm-and-conv/
- with torch.no_grad():
- # init
- fusedconv = torch.nn.Conv2d(conv.in_channels,
- conv.out_channels,
- kernel_size=conv.kernel_size,
- stride=conv.stride,
- padding=conv.padding,
- bias=True)
-
- # 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_(b_conv + b_bn)
-
- return fusedconv
-
-
- def model_info(model, report='summary'):
- # 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
- if report is 'full':
- 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()))
- print('Model Summary: %g layers, %g parameters, %g gradients' % (len(list(model.parameters())), n_p, n_g))
|