car-detection-bayes/utils/torch_utils.py

61 lines
2.1 KiB
Python
Raw Normal View History

import torch
def init_seeds(seed=0):
torch.manual_seed(seed)
2019-02-25 12:47:51 +00:00
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
2019-06-24 11:43:17 +00:00
# torch.backends.cudnn.deterministic = True # https://pytorch.org/docs/stable/notes/randomness.html
def select_device(force_cpu=False):
2019-04-08 13:41:14 +00:00
cuda = False if force_cpu else torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
2019-02-16 13:33:52 +00:00
2019-04-08 13:41:14 +00:00
if not cuda:
print('Using CPU')
if cuda:
2019-06-24 11:51:54 +00:00
torch.backends.cudnn.benchmark = True # set False for reproducible results
2019-04-08 13:41:14 +00:00
c = 1024 ** 2 # bytes to MB
ng = torch.cuda.device_count()
x = [torch.cuda.get_device_properties(i) for i in range(ng)]
print("Using CUDA device0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
(x[0].name, x[0].total_memory / c))
if ng > 0:
# torch.cuda.set_device(0) # OPTIONAL: Set GPU ID
for i in range(1, ng):
print(" device%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
(i, x[i].name, x[i].total_memory / c))
2019-02-16 13:33:52 +00:00
2019-05-03 16:14:16 +00:00
print('') # skip a line
return device
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
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