67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
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
|