summary | shortlog | log | commit | commitdiff | tree
raw | patch | inline | side by side (parent: d53c246)
raw | patch | inline | side by side (parent: d53c246)
author | Manu Mathew <a0393608@ti.com> | |
Mon, 11 May 2020 20:51:27 +0000 (02:21 +0530) | ||
committer | Manu Mathew <a0393608@ti.com> | |
Mon, 11 May 2020 20:53:19 +0000 (02:23 +0530) |
diff --git a/modules/pytorch_jacinto_ai/vision/models/mobilenetv1.py b/modules/pytorch_jacinto_ai/vision/models/mobilenetv1.py
index c9ee147a247ca731d424bc8053f784f4ec0637ad..0139e156ec84d5e1b4daa9149005a2b7263034ad 100644 (file)
# building first layer
output_channels = int(self.model_config.layer_setting[0][1] * width_mult)
- output_channels = make_divisible_by8(output_channels) if model_config.en_make_divisible_by8 else output_channels
+ output_channels = xnn.utils.make_divisible_by8(output_channels) if model_config.en_make_divisible_by8 else output_channels
features = [xnn.layers.ConvNormAct2d(3, output_channels, kernel_size=kernel_size, stride=s0, activation=activation)]
channels = output_channels
# building inverted residual blocks
for t, c, n, s in self.model_config.layer_setting[1:]:
- output_channels = make_divisible_by8(c * width_mult) if model_config.en_make_divisible_by8 else int(c * width_mult)
+ output_channels = xnn.utils.make_divisible_by8(c * width_mult) if model_config.en_make_divisible_by8 else int(c * width_mult)
for i in range(n):
stride = s if i == 0 else 1
block = BlockBuilder(channels, output_channels, stride=stride, kernel_size=kernel_size, activation=(activation,activation))
diff --git a/modules/pytorch_jacinto_ai/vision/models/mobilenetv2.py b/modules/pytorch_jacinto_ai/vision/models/mobilenetv2.py
index 89b59e2331ce24bbf64053394891223a0b50db81..3cb392c3580c94dcc3e6519dedccbe6e8a83913e 100644 (file)
kernel_size = self.model_config.kernel_size
# building first layer
- output_channels = make_divisible_by8(self.model_config.layer_setting[0][1] * width_mult)
+ output_channels = xnn.utils.make_divisible_by8(self.model_config.layer_setting[0][1] * width_mult)
features = [xnn.layers.ConvNormAct2d(model_config.input_channels, output_channels, kernel_size=kernel_size, stride=s0, activation=activation)]
channels = output_channels
# building inverted residual blocks
for t, c, n, s in self.model_config.layer_setting[1:-1]:
- output_channels = make_divisible_by8(c * width_mult)
+ output_channels = xnn.utils.make_divisible_by8(c * width_mult)
for i in range(n):
stride = s if i == 0 else 1
block = BlockBuilder(channels, output_channels, stride=stride, kernel_size=kernel_size, activation=activation, expand_ratio=t, linear_dw=linear_dw)
# building classifier
if self.model_config.num_classes != None:
- output_channels = make_divisible_by8(self.model_config.layer_setting[-1][1] * width_mult)
+ output_channels = xnn.utils.make_divisible_by8(self.model_config.layer_setting[-1][1] * width_mult)
features.append(xnn.layers.ConvNormAct2d(channels, output_channels, kernel_size=1, activation=activation))
channels = output_channels
self.classifier = torch.nn.Sequential(
diff --git a/modules/pytorch_jacinto_ai/vision/models/utils.py b/modules/pytorch_jacinto_ai/vision/models/utils.py
index bf886cfdbe76e422dba27a239e9719e15b7bc48b..c2f7385660281f0c4b52c576e3604589c6dfd6e9 100644 (file)
from torch.utils.model_zoo import load_url as load_state_dict_from_url
-def make_divisible(v, divisor, min_value=None):
- """
- This function is taken from the original tf repo.
- It ensures that all layers have a channel number that is divisible by 8
- It can be seen here:
- https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
- :param v:
- :param divisor:
- :param min_value:
- :return:
- """
- if min_value is None:
- min_value = divisor
- new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
- # Make sure that round down does not go down by more than 10%.
- if new_v < 0.9 * v:
- new_v += divisor
- return int(new_v)
-
-def make_divisible_by8(v):
- return make_divisible(v, 8)
diff --git a/modules/pytorch_jacinto_ai/xnn/quantize/quant_base_module.py b/modules/pytorch_jacinto_ai/xnn/quantize/quant_base_module.py
index 6e50118d1c3ea9e9bc8e48cace6fedfc2d25fddb..c53c3e4a7e61cc88f0f13c5d05ebcca29ebc3049 100644 (file)
# base module to be use for all quantization modules
class QuantBaseModule(QuantGraphModule):
- def __init__(self, module, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False,
- histogram_range=True, bias_calibration=False, constrain_weights=False, dummy_input=None,
+ def __init__(self, module, dummy_input, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False,
+ histogram_range=True, bias_calibration=False, constrain_weights=False,
model_surgery_quantize=False):
super().__init__(module)
self.bitwidth_weights = bitwidth_weights
diff --git a/modules/pytorch_jacinto_ai/xnn/quantize/quant_calib_module.py b/modules/pytorch_jacinto_ai/xnn/quantize/quant_calib_module.py
index 2ab33d8ee86e8926d25aa0fee4aaf8a881c816e9..00227a0918a9b2c8ca06b1b2e51234503866ab1e 100644 (file)
###########################################################
class QuantCalibrateModule(QuantTrainModule):
- def __init__(self, module, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False,
- histogram_range=True, bias_calibration=True, constrain_weights=True, dummy_input=None, lr_calib=0.05):
+ def __init__(self, module, dummy_input, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False,
+ histogram_range=True, bias_calibration=True, constrain_weights=True, lr_calib=0.05):
self.bias_calibration = bias_calibration
self.weights_calibration = False
self.lr_calib = lr_calib
self.calibrate_repeats = 1
self.quantize_enable = True
self.update_range = True
- super().__init__(module, bitwidth_weights=bitwidth_weights, bitwidth_activations=bitwidth_activations,
+ super().__init__(module, dummy_input=dummy_input, bitwidth_weights=bitwidth_weights, bitwidth_activations=bitwidth_activations,
per_channel_q=per_channel_q, histogram_range=histogram_range, bias_calibration=bias_calibration,
- constrain_weights=constrain_weights, dummy_input=dummy_input)
+ constrain_weights=constrain_weights)
def forward(self, inputs):
diff --git a/modules/pytorch_jacinto_ai/xnn/quantize/quant_test_module.py b/modules/pytorch_jacinto_ai/xnn/quantize/quant_test_module.py
index 7166caade59c444511edbe57ec147cb33eba2fab..557a504dd90350a161b548d1f2d1bdfab9f8ab73 100644 (file)
class QuantTestModule(QuantBaseModule):
- def __init__(self, module, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False, histogram_range=True,
- range_calibration_online=False, bias_calibration=False, dummy_input=None, model_surgery_quantize=False):
- super().__init__(module, bitwidth_weights=bitwidth_weights, bitwidth_activations=bitwidth_activations,
- per_channel_q=per_channel_q, histogram_range=histogram_range, bias_calibration=bias_calibration,
- constrain_weights=False, dummy_input=dummy_input, model_surgery_quantize=model_surgery_quantize)
+ def __init__(self, module, dummy_input, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False, histogram_range=True,
+ range_calibration_online=False, model_surgery_quantize=False):
+ super().__init__(module, dummy_input=dummy_input, bitwidth_weights=bitwidth_weights, bitwidth_activations=bitwidth_activations,
+ per_channel_q=per_channel_q, histogram_range=histogram_range, bias_calibration=False,
+ constrain_weights=False, model_surgery_quantize=model_surgery_quantize)
# use power2_weights for now
self.power2_weights = True
# whether to do online adjustment of calibration using previous frame range
# calibration does not need gradients
with torch.no_grad():
- if self.bias_calibration:
- # bias calibration - not implemented in this base class
- self.calibrate_bias(inputs)
- else:
- # quantize
- outputs = self.forward_quantize(inputs)
- #
+ # quantize
+ outputs = self.forward_quantize(inputs)
# start and new frame, copy the qparams for previous frame of inference
self.get_state().qparams_prev = self.copy_qparams(self.get_state().qparams, inputs)
# return
self.idx_large_mse_for_act += 1
-# class QuantCalibrateTestModule(QuantTestModule):
-# def __init__(self, module, bitwidth_weights, bitwidth_activations, per_channel_q, power2_weights=True, histogram_range=True,
-# range_calibration_online=False, bias_calibration=False, model_surgery_quantize=False, dummy_input=None):
-# super().__init__(module, bitwidth_weights, bitwidth_activations, per_channel_q, power2_weights, histogram_range, range_calibration_online,
-# bias_calibration=bias_calibration, model_surgery_quantize=model_surgery_quantize, dummy_input=dummy_input)
-# self.bias_calibration = bias_calibration
-# # bias correction update factor
-# self.max_calibration_repeats = max(30//self.range_calibration_offline_iters,1)
-# self.bias_calibration_factor = 1.0/(self.range_calibration_offline_iters*self.max_calibration_repeats)
-# self.constrain_weights_enabled = bias_calibration #and (not per_channel_q)
-# self.adjust_bias_now = False
-# # TODO: it may be better to start bias adjustment after the activation range update has stabilized
-# self.bias_calibration_start = 0
-# #
-#
-#
-# def calibrate_bias(self, inputs):
-# is_bias_calibration = (self.bias_calibration and self.iter_in_epoch>=self.bias_calibration_start and \
-# self.iter_in_epoch < self.range_calibration_offline_iters)
-# if not is_bias_calibration:
-# return
-# #
-# # do not update activation range during bias adjustment
-# # backup weghts in the first iteration
-# if (self.iter_in_epoch == self.bias_calibration_start):
-# # backup original weights
-# self._backup_weights_orig()
-# # quantize weights
-# self.forward_quantize(inputs)
-# # backup quantized weights
-# self._backup_weights_quant()
-# #
-# self.forward_compute_oputput_stats(inputs)
-#
-# # bias adjustment
-# self.forward_adjust_bias(inputs)
-#
-# if (self.iter_in_epoch == self.range_calibration_offline_iters):
-# self._remove_backups()
-# #
-# #
-#
-#
-# def apply_constrain_weights(self, module):
-# if not self.constrain_weights_enabled:
-# return
-# #
-# constrained_weight = quant_utils.constrain_weight(module.weight.data)
-# module.weight.data.copy_(constrained_weight.data)
-# #
-#
-#
-# def forward_compute_oputput_stats(self, inputs):
-# # restore
-# self._restore_weights_orig()
-# # forward with hook
-# self.add_call_hook(self.module, self._forward_compute_oputput_stats_hook)
-# _ = self.module(inputs)
-# #
-# def _forward_compute_oputput_stats_hook(self, op, *inputs_orig):
-# outputs = op.__forward_orig__(*inputs_orig)
-# # bias adjustment
-# bias_calibration_op = hasattr(op, 'bias') and (op.bias is not None)
-# if bias_calibration_op:
-# output = outputs[0] if isinstance(outputs, (list,tuple)) else outputs
-# channels = output.size(1)
-# op.__output_mean__orig__ = output.mean(0).view(channels,-1).mean(-1)
-# #
-# return outputs
-# #
-#
-#
-# def forward_adjust_bias(self, input):
-# # restore
-# self._restore_weights_quant()
-# self.adjust_bias_now = True
-# for rpt_idx in range(self.max_calibration_repeats):
-# # actual bias adjustment
-# output = super().forward_quantize(input)
-# #
-# self.adjust_bias_now = False
-# # backup quantized weights
-# self._backup_weights_quant()
-# return output
-# #
-#
-#
-# # called during forward_quantize()
-# def process_outputs(self, op, inputs, outputs):
-# super().process_outputs(op, inputs, outputs)
-# # do adjustment of bias
-# bias_calibration_op = self.adjust_bias_now and hasattr(op, 'bias') and (op.bias is not None)
-# if bias_calibration_op:
-# channels = outputs.size(1)
-# output_mean = outputs.mean(0).view(channels,-1).mean(-1)
-# output_delta = op.__output_mean__orig__ - output_mean
-# output_delta_scaled = output_delta * self.bias_calibration_factor
-# # update the bias as well as the backup copy
-# op.__bias__quant__ += (output_delta_scaled)
-# op.bias.data += (output_delta_scaled)
-# #TODO: is this needed?
-# #outputs.add_(output_delta_scaled.view(1,-1,1,1))
-# #
-# #
-#
diff --git a/modules/pytorch_jacinto_ai/xnn/quantize/quant_train_module.py b/modules/pytorch_jacinto_ai/xnn/quantize/quant_train_module.py
index 2344b653ed05a8ab7c2cc5739409ad59a8eab720..85319f0212dde5d0d3b925921eb2982764db5aa9 100644 (file)
###########################################################
class QuantTrainModule(QuantBaseModule):
- def __init__(self, module, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False, histogram_range=True,
- bias_calibration=False, constrain_weights=True, dummy_input=None):
- super().__init__(module, bitwidth_weights=bitwidth_weights, bitwidth_activations=bitwidth_activations,
+ def __init__(self, module, dummy_input, bitwidth_weights=8, bitwidth_activations=8, per_channel_q=False,
+ histogram_range=True, bias_calibration=False, constrain_weights=True):
+ super().__init__(module, dummy_input=dummy_input, bitwidth_weights=bitwidth_weights, bitwidth_activations=bitwidth_activations,
per_channel_q=per_channel_q, histogram_range=histogram_range, bias_calibration=bias_calibration,
- constrain_weights=constrain_weights, dummy_input=dummy_input, model_surgery_quantize=True)
+ constrain_weights=constrain_weights, model_surgery_quantize=True)
# range shrink - 0.0 indicates no shrink
percentile_range_shrink = (layers.PAct2.PACT2_RANGE_SHRINK if histogram_range else 0.0)
# set attributes to all modules - can control the behaviour from here
diff --git a/modules/pytorch_jacinto_ai/xnn/utils/__init__.py b/modules/pytorch_jacinto_ai/xnn/utils/__init__.py
index b4bc88118acc1c483f3f7235f2cffbf672b6cf66..2532bfe9854b47c8e1415f4d10485eb9a5ce9231 100644 (file)
from .utils_data import *
from .load_weights import *
from .tensor_utils import *
-#from .tensor_utils_internal import *
from .logger import *
from .utils_hist import *
from .attr_dict import *
diff --git a/modules/pytorch_jacinto_ai/xnn/utils/count_flops.py b/modules/pytorch_jacinto_ai/xnn/utils/count_flops.py
index 144fd2e547701d97321a65690c76a79bad33c95e..bf63abbcfd5a7e3c795c33d1019c2e0aac4848be 100644 (file)
import torch
from . import module_utils
-def forward_count_flops(module, inp):
- _add_hook(module, _count_flops_func)
- _ = module(inp)
+def forward_get_complexity(model, inp):
+ num_flops = forward_count_flops(model, inp)
+ num_params = count_params(model)
+ return num_flops, num_params
+
+
+def forward_count_flops(model, inp):
+ _add_hook(model, _count_flops_func)
+ _ = model(inp)
num_flops = 0
- for m in module.modules():
+ for m in model.modules():
num_flops += m.__num_flops__
#
- _remove_hook(module)
+ _remove_hook(model)
return num_flops
+def count_params(model):
+ layer_params = [p.numel() for p in model.parameters if p.requires_grad]
+ num_params = sum(layer_params)
+ return num_params
+
+
def _count_flops_func(m, inp, out):
# trained calibration/quantization can do model surgery and return extra outputs - ignroe them
if isinstance(out, (list,tuple)):
diff --git a/modules/pytorch_jacinto_ai/xnn/utils/tensor_utils.py b/modules/pytorch_jacinto_ai/xnn/utils/tensor_utils.py
index 3de9aa13ad4d01005bc4c9ef1c77b1b4763e7fe0..70c90c54118edeeb00f45b13a9c347d05cc23933 100644 (file)
import torch
import scipy
import warnings
+import cv2
from ..layers import functional
-from . import util_functions
+from . import image_utils
###############################################################
return new_mn_scaled, new_mx_scaled
-##################################################################
-def shrink_tensor(tensor_in, range_shrink):
- weight_low, weight_high = extrema(tensor_in, percentile_range_shrink=range_shrink)
- weight = torch.min(torch.max(tensor_in, weight_low), weight_high)
- return weight
-
-
##################################################################
def check_sizes(input, input_name, expected):
condition = [input.ndimension() == len(expected)]
assert(all(condition)), "wrong size for {}, expected {}, got {}".format(input_name, 'x'.join(expected), list(input.size()))
+###########################################################################
+def tensor2img(tensor, adjust_range=True, min_value = None, max_value=None):
+ if tensor.ndimension() < 3:
+ tensor = tensor.unsqueeze(0)
+ if tensor.ndimension() < 4:
+ tensor = tensor.unsqueeze(0)
+ if min_value is None:
+ min_value = tensor.min()
+ if max_value is None:
+ max_value = tensor.max()
+ range = max_value-min_value
+ array = (255*(tensor - min_value)/range).clamp(0,255) if adjust_range else tensor
+ if array.size(1) >= 3:
+ img = torch.stack((array[0,0], array[0,1], array[0,2]), dim=2)
+ else:
+ img = array[0,0]
+ return img.cpu().data.numpy().astype(np.uint8)
+
+
+def flow2rgb(flow_map, max_value):
+ global args
+ _, h, w = flow_map.shape
+ #flow_map[:,(flow_map[0] == 0) & (flow_map[1] == 0)] = float('nan')
+ rgb_map = np.ones((h,w,3)).astype(np.float32)
+ if max_value is not None:
+ normalized_flow_map = flow_map / max_value
+ else:
+ normalized_flow_map = flow_map / (np.abs(flow_map).max())
+ rgb_map[:,:,0] += normalized_flow_map[0]
+ rgb_map[:,:,1] -= 0.5*(normalized_flow_map[0] + normalized_flow_map[1])
+ rgb_map[:,:,2] += normalized_flow_map[1]
+ return rgb_map.clip(0,1)
+
+
+def flow2hsv(flow_map, max_value=128, scale_fact=8, confidence=False):
+ global args
+ _, h, w = flow_map.shape
+ hsv = np.zeros((h, w, 3)).astype(np.float32)
+
+ mag = np.sqrt(flow_map[0]**2 + flow_map[1]**2)
+ phase = np.arctan2(flow_map[1], flow_map[0])
+ phase = np.mod(phase/(2*np.pi), 1)
+
+ hsv[:, :, 0] = phase*360
+ hsv[:, :, 1] = (mag*scale_fact/max_value).clip(0, 1)
+ hsv[:, :, 2] = (scale_fact - hsv[:, :, 1]).clip(0, 1)
+ rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
+ if confidence:
+ return rgb * flow_map[2] > 128
+ else:
+ return rgb
+
+
+def tensor2array(tensor, max_value=255.0, colormap='rainbow', input_blend=None):
+ max_value = float(tensor.max()) if max_value is None else max_value
+
+ if tensor.ndimension() == 2 or tensor.size(0) == 1:
+ try:
+ import cv2
+ if cv2.__version__.startswith('2') :
+ color_cvt = cv2.cv.CV_BGR2RGB
+ else: # 3.x,4,x
+ color_cvt = cv2.COLOR_BGR2RGB
+ #
+ if colormap == 'rainbow':
+ colormap = cv2.COLORMAP_RAINBOW
+ elif colormap == 'magma': # >=3.4.8
+ colormap = cv2.COLORMAP_MAGMA
+ elif colormap == 'bone':
+ colormap = cv2.COLORMAP_BONE
+ elif colormap == 'plasma': # >=4.1
+ colormap = cv2.COLORMAP_PLASMA
+ elif colormap == 'turbo': # >=4.1.2
+ colormap = cv2.COLORMAP_TURBO
+ #
+ array = (255.0*tensor.squeeze().numpy()/max_value).clip(0, 255).astype(np.uint8)
+ colored_array = cv2.applyColorMap(array, colormap)
+ array = cv2.cvtColor(colored_array, color_cvt).astype(np.float32) / 255.0
+ except ImportError:
+ if tensor.ndimension() == 2:
+ tensor.unsqueeze_(2)
+ #
+ array = (tensor.expand(tensor.size(0), tensor.size(1), 3).numpy()/max_value).clip(0,1)
+ elif tensor.ndimension() == 3:
+ assert(tensor.size(0) == 3)
+ array = 0.5 + tensor.numpy().transpose(1, 2, 0)*0.5
+ #
+ if input_blend is not None:
+ array = image_utils.chroma_blend(input_blend, array)
+ #
+ return array
+
+
+def tensor2img(tensor, max_value=63535):
+ array = (63535*tensor.numpy()/max_value).clip(0, 63535).astype(np.uint16)
+ if tensor.ndimension() == 3:
+ assert (array.size(0) == 3)
+ array = array.transpose(1, 2, 0)
+ return array
+
+
##################################################################
def inverse_warp_flow(img, flow, padding_mode='zeros'):
"""
def debug_dump_tensor(tensor, image_name, adjust_range=True):
- img = util_functions.tensor2img(tensor, adjust_range=adjust_range)
+ img = tensor2img(tensor, adjust_range=adjust_range)
scipy.misc.imsave(image_name, img)
diff --git a/modules/pytorch_jacinto_ai/xnn/utils/util_functions.py b/modules/pytorch_jacinto_ai/xnn/utils/util_functions.py
index fb74244c8afece424e85cc72b2251afae98f4714..d7c1228ba4e9ab99472bd87227592c34e5e0f8ee 100644 (file)
v[index] = str2bool(args)
return v
+
+#########################################################################
+def make_divisible(value, factor, min_value=None):
+ """
+ Inspired by https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+ """
+ min_value = factor if min_value is None else min_value
+ round_factor = factor/2
+ quotient = int(value + round_factor) // factor
+ value_multiple = max(quotient * factor, min_value)
+ # make sure that the change is contained
+ if value_multiple < 0.9*value:
+ value_multiple = value_multiple + factor
+ #
+ return int(value_multiple)
+
+
+def make_divisible_by8(v):
+ return make_divisible(v, 8)
+
+
+#########################################################################
def recursive_glob(rootdir='.', suffix=''):
"""Performs recursive glob with given suffix and rootdir
:param rootdir is the root directory
for filename in filenames if filename.endswith(suffix)]
-
-#########################################################################
-def resize(input, size, sparse, flow_scale=False):
- size_in = input.size()[2:]
- scale_fact = [so/si for si, so in zip(size_in, size)]
- if scale_fact[0] == 1 and scale_fact[1] == 1: #{
- output = input
- elif scale_fact[0] > 1 and scale_fact[1] > 1:
- output = torch.nn.functional.interpolate(input, size=size, mode='bilinear')
- elif sparse:
- output = sparse_resample(input, size)
- else:
- output = torch.nn.functional.adaptive_avg_pool2d(input, size)
- #}
- if flow_scale:
- output[:,0,...] = output[:,0,...] * scale_fact[1] #x_flow
- output[:,1,...] = output[:,1,...] * scale_fact[0] #y_flow
-
- return output
-
-
-def sparse_resample(input, size, step_sampling=True):
- size_in = input.size()[2:]
- scale = [so/si for si, so in zip(size_in, size)]
- if scale[0] == 1 and scale[1] == 1:
- output = input
- elif scale[0] > 1 and scale[1] > 1:
- output = torch.nn.functional.interpolate(input, size=size, mode='nearest')
- elif step_sampling:
- # step sampling acts like nearest neighbor resampling,
- # not sure if this is the right thing for multi scale learning
- size_new = [(so * si) // so for si, so in zip(size_in, size)]
- if size_new[0] != size_in[0] or size_new[1] != size_in[1]:
- input_new = sparse_max_pool_(input, size_new)
- else:
- input_new = input
- output = sparse_resample_step_(input_new, size)
- #debug_dump_tensor(input, 'orig_input.png', adjust_range=False)
- #debug_dump_tensor(output, 'resized_target.png', adjust_range=False)
- else:
- output = sparse_max_pool_(input, size)
- return output
-
-
-def sparse_max_pool_(input, size):
- positive = (input > 0).float()
- negative = (input < 0).float()
- output = torch.nn.functional.adaptive_max_pool2d(input * positive, size) - \
- torch.nn.functional.adaptive_max_pool2d(-input * negative, size)
- return output
-
-
-def sparse_resample_step_(input, size):
- size_in = input.size()[2:]
- step = [si//so for si, so in zip(size_in, size)]
- output = input[...,::step[0],::step[1]]
- return output
-
-
-def upsample_flow(flow, pred_type, mode, size=None, scale_factor=None):
- assert (size is not None or scale_factor is not None), \
- 'either scale_factor or size should be specified'
- assert (size is None or scale_factor is None), \
- 'both scale_factor and size should not be specified'
- flow_up = torch.nn.functional.interpolate(flow, size=size, scale_factor=scale_factor, mode=mode)
-
- if 'flow' in pred_type:
- scale_factor = torch.autograd.Variable(torch.ones(1, flow_up.size(1), 1, 1), requires_grad=False).cuda()
- scale_factor[:, 0,...] = (float(flow_up.size(3)) / float(flow.size(3))) # hor flow
- scale_factor[:, 1,...] = (float(flow_up.size(2)) / float(flow.size(2))) # ver flow
- flow_up = flow_up * scale_factor
-
- return flow_up
-
-
-def upsample_flow2(flow, pred_type, mode):
- return upsample_flow(flow, pred_type, mode, scale_factor=2)
-
-class UpsampleFlow(torch.nn.Module):
- def __init__(self, pred_type, mode, size, scale_factor):
- self.size = size
- self.scale_factor = scale_factor
- self.pred_type = pred_type
- self.mode = mode
- super().__init__()
-
- def forward(self, flow):
- return upsample_flow(flow, self.pred_type, self.mode, self.size, self.scale_factor)
-
-
-class UpsampleFlow2(torch.nn.Module):
- def __init__(self, pred_type, mode):
- self.pred_type = pred_type
- self.mode = mode
- super().__init__()
-
- def forward(self, flow):
- return upsample_flow(flow, self.pred_type, self.mode, size=None, scale_factor=2)
-
-
-def downsample_flow(flow, pred_type, mode, size=None, scale_factor=None):
- assert (size is not None or scale_factor is not None), \
- 'either scale_factor or size should be specified'
- assert (size is None or scale_factor is None), \
- 'both scale_factor and size should not be specified'
- if scale_factor is not None:
- size = (flow.size(2)//scale_factor, flow.size(3)//scale_factor)
-
- flow_down = torch.nn.functional.adaptive_avg_pool2d(flow, output_size=size)
- if 'flow' in pred_type:
- scale_factor = torch.autograd.Variable(torch.ones(1, flow_down.size(1), 1, 1), requires_grad=False).cuda()
- scale_factor[:, 0,...] = (float(flow_down.size(3)) / float(flow.size(3))) # hor flow
- scale_factor[:, 1,...] = (float(flow_down.size(2)) / float(flow.size(2))) # ver flow
- flow_down = flow_down * scale_factor
- return flow_down
-
-
-def downsample_flow2(flow, pred_type, mode):
- return downsample_flow(flow, pred_type, mode, scale_factor=2)
-
-
-###########################################################################
-def tensor2img(tensor, adjust_range=True, min_value = None, max_value=None):
- if tensor.ndimension() < 3:
- tensor = tensor.unsqueeze(0)
- if tensor.ndimension() < 4:
- tensor = tensor.unsqueeze(0)
- if min_value is None:
- min_value = tensor.min()
- if max_value is None:
- max_value = tensor.max()
- range = max_value-min_value
- array = (255*(tensor - min_value)/range).clamp(0,255) if adjust_range else tensor
- if array.size(1) >= 3:
- img = torch.stack((array[0,0], array[0,1], array[0,2]), dim=2)
- else:
- img = array[0,0]
- return img.cpu().data.numpy().astype(np.uint8)
-
-
-def flow2rgb(flow_map, max_value):
- global args
- _, h, w = flow_map.shape
- #flow_map[:,(flow_map[0] == 0) & (flow_map[1] == 0)] = float('nan')
- rgb_map = np.ones((h,w,3)).astype(np.float32)
- if max_value is not None:
- normalized_flow_map = flow_map / max_value
- else:
- normalized_flow_map = flow_map / (np.abs(flow_map).max())
- rgb_map[:,:,0] += normalized_flow_map[0]
- rgb_map[:,:,1] -= 0.5*(normalized_flow_map[0] + normalized_flow_map[1])
- rgb_map[:,:,2] += normalized_flow_map[1]
- return rgb_map.clip(0,1)
-
-
-def flow2hsv(flow_map, max_value=128, scale_fact=8, confidence=False):
- global args
- _, h, w = flow_map.shape
- hsv = np.zeros((h, w, 3)).astype(np.float32)
-
- mag = np.sqrt(flow_map[0]**2 + flow_map[1]**2)
- phase = np.arctan2(flow_map[1], flow_map[0])
- phase = np.mod(phase/(2*np.pi), 1)
-
- hsv[:, :, 0] = phase*360
- hsv[:, :, 1] = (mag*scale_fact/max_value).clip(0, 1)
- hsv[:, :, 2] = (scale_fact - hsv[:, :, 1]).clip(0, 1)
- rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
- if confidence:
- return rgb * flow_map[2] > 128
- else:
- return rgb
-
-
-
-def tensor2array(tensor, max_value=255.0, colormap='rainbow', input_blend=None):
- max_value = float(tensor.max()) if max_value is None else max_value
-
- if tensor.ndimension() == 2 or tensor.size(0) == 1:
- try:
- import cv2
- if cv2.__version__.startswith('2') :
- color_cvt = cv2.cv.CV_BGR2RGB
- else: # 3.x,4,x
- color_cvt = cv2.COLOR_BGR2RGB
- #
- if colormap == 'rainbow':
- colormap = cv2.COLORMAP_RAINBOW
- elif colormap == 'magma': # >=3.4.8
- colormap = cv2.COLORMAP_MAGMA
- elif colormap == 'bone':
- colormap = cv2.COLORMAP_BONE
- elif colormap == 'plasma': # >=4.1
- colormap = cv2.COLORMAP_PLASMA
- elif colormap == 'turbo': # >=4.1.2
- colormap = cv2.COLORMAP_TURBO
- #
- array = (255.0*tensor.squeeze().numpy()/max_value).clip(0, 255).astype(np.uint8)
- colored_array = cv2.applyColorMap(array, colormap)
- array = cv2.cvtColor(colored_array, color_cvt).astype(np.float32) / 255.0
- except ImportError:
- if tensor.ndimension() == 2:
- tensor.unsqueeze_(2)
- #
- array = (tensor.expand(tensor.size(0), tensor.size(1), 3).numpy()/max_value).clip(0,1)
- elif tensor.ndimension() == 3:
- assert(tensor.size(0) == 3)
- array = 0.5 + tensor.numpy().transpose(1, 2, 0)*0.5
- #
- if input_blend is not None:
- array = chroma_blend(input_blend, array)
- #
- return array
-
-
-def tensor2img(tensor, max_value=63535):
- array = (63535*tensor.numpy()/max_value).clip(0, 63535).astype(np.uint16)
- if tensor.ndimension() == 3:
- assert (array.size(0) == 3)
- array = array.transpose(1, 2, 0)
- return array
-
-
-def compute_errors(gt, pred, sparse, max_depth, crop_gt=True):
- #default is a very large value, which effectively does nothing
- max_depth = max_depth if max_depth is not None else 10000.0
- mask = generate_mask(gt, 0, max_depth, crop_gt=crop_gt)
- pred = pred[mask]
- gt = gt[mask]
-
- if sparse:
- pred = pred[gt!=0]
- gt = gt[gt!=0]
- #if max_depth is not None:
- # pred = pred[gt<max_depth]
- # gt = gt[gt<max_depth]
-
- thresh = np.maximum((gt / pred), (pred / gt))
- a1 = (thresh < 1.25 ).mean()
- a2 = (thresh < 1.25 ** 2).mean()
- a3 = (thresh < 1.25 ** 3).mean()
-
- rmse = (gt - pred) ** 2
- rmse = np.sqrt(rmse.mean())
-
- rmse_log = (np.log(gt) - np.log(pred)) ** 2
- rmse_log = np.sqrt(rmse_log.mean())
-
- abs_rel = np.mean(np.abs(gt - pred) / gt)
- sq_rel = np.mean(((gt - pred)**2) / gt)
- return abs_rel, sq_rel, rmse, rmse_log, a1, a2, a3
-
-
-def generate_mask(gt_depth, min_depth, max_depth, crop_gt=True):
- mask = np.logical_and(gt_depth > min_depth,
- gt_depth < max_depth)
- if crop_gt:
- # crop used by Garg ECCV16 to reprocude Eigen NIPS14 results
- # if used on gt_size 370x1224 produces a crop of [-218, -3, 44, 1180]
- gt_height, gt_width = gt_depth.shape
- crop = np.array([0.40810811 * gt_height, 0.99189189 * gt_height,
- 0.03594771 * gt_width, 0.96405229 * gt_width]).astype(np.int32)
-
- crop_mask = np.zeros(mask.shape)
- crop_mask[crop[0]:crop[1],crop[2]:crop[3]] = 1
- mask = np.logical_and(mask, crop_mask)
- return mask
-
-
###############################################################
def get_shape_with_stride(in_shape, stride):
shape_s = [in_shape[0],in_shape[1],in_shape[2]//stride,in_shape[3]//stride]