def _prepare_qdq_export_quantization(self, x: torch.Tensor):
x, level_high, level_low, input_low, input_high = self._prepare_export_quantization(
x)
with no_jit_trace():
y_scale, y_zero_point = get_scale_zp_from_input_low_input_high(
level_low, level_high, input_low, input_high)
return x, y_scale, y_zero_point
def forward(self, **params):
new_params = []
for param_name, param_value in params.items():
# In case of None weight (or bias) mask shouldn't be applied
if param_value is None:
new_params.append(param_value)
continue
# For weights self.mask_applying_dim should be used, for bias dim=0
dim = 0 if param_name == 'bias' else self.mask_applying_dim
if is_tracing_state():
with no_jit_trace():
new_params.append(
inplace_apply_filter_binary_mask(
self.binary_filter_pruning_mask,
param_value,
node_name_for_logging=self.node_name,
dim=dim))
else:
new_params.append(
apply_filter_binary_mask(
self.binary_filter_pruning_mask,
param_value,
node_name_for_logging=self.node_name,
dim=dim))
return new_params
def _prepare_export_quantization(self, x: torch.Tensor):
with no_jit_trace():
input_low, input_high = self._get_input_low_input_high(
self.scale, self.level_low, self.level_high, self.eps)
level_low = self.level_low
level_high = self.level_high
if self._half_range:
x = torch.min(torch.max(x, input_low), input_high)
level_low = 2 * self.level_low
level_high = 2 * self.level_high + 1
input_low, input_high = self._get_input_low_input_high(
level_high / self.level_high * self.scale, level_low,
level_high, self.eps)
return x, level_high, level_low, input_low, input_high
def _conv_forward_proxy(self, input_, weight, bias, padding_value):
with no_jit_trace():
padding_val = padding_value.item()
self.get_padding_value_ref().data.fill_(padding_val)
if self.padding_mode != 'zeros':
return F.conv2d(F.pad(input_, self._reversed_padding_repeated_twice, mode=self.padding_mode,
value=padding_val),
weight, bias, self.stride,
(0, 0), self.dilation, self.groups)
if padding_val == 0:
return F.conv2d(input_, weight, bias, self.stride, self.padding, self.dilation, self.groups)
return F.conv2d(F.pad(input_, self._reversed_padding_repeated_twice, value=padding_val),
weight, bias, self.stride,
(0, 0), self.dilation, self.groups)
def forward(self, weight):
if is_tracing_state():
with no_jit_trace():
return weight.mul_(self.binary_mask)
tmp_tensor = self._calc_training_binary_mask(weight)
return apply_binary_mask_impl(tmp_tensor, weight)
def signed(self):
with no_jit_trace():
return self.signed_tensor.item() == 1
def _prepare_fq_export_quantization(self, x: torch.Tensor):
x, level_high, level_low, input_low, input_high = self._prepare_export_quantization(
x)
with no_jit_trace():
levels = level_high - level_low + 1
return x, levels, input_low, input_high
def num_bits(self):
with no_jit_trace():
return self._num_bits.item()
def is_enabled_quantization(self):
with no_jit_trace():
return self.enabled[0].item() == 1
