def test_input_weight_equalization_weights_bias(self): """ After applying the equalization functions check if the weights and biases are as expected """ tests = [SingleLayerLinearModel, TwoLayerLinearModel, SingleLayerFunctionalLinearModel, TwoLayerFunctionalLinearModel] x = torch.rand((5, 5)) for M in tests: m = M().eval() exp_eq_scales = self.get_expected_eq_scales(m, x.detach().numpy()) exp_weights, exp_bias = self.get_expected_weights_bias(m, x.detach().numpy(), exp_eq_scales) prepared = prepare_fx(m, specific_qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) prepared(x) convert_ref = _convert_equalization_ref(prepared) convert_ref(x) modules = dict(convert_ref.named_modules(remove_duplicate=False)) counter = 0 for node in convert_ref.graph.nodes: if node.op == 'call_module' and isinstance(modules[str(node.target)], nn.Linear): self.assertEqual(modules[str(node.target)].weight, exp_weights[counter]) self.assertEqual(modules[str(node.target)].bias, exp_bias[counter]) counter += 1
def test_input_weight_equalization_equalization_scales(self): """ After applying the equalization functions, check if the equalization scales are the expected values """ tests = [ SingleLayerLinearModel, TwoLayerLinearModel, SingleLayerFunctionalLinearModel, TwoLayerFunctionalLinearModel ] x = torch.rand((5, 5)) for M in tests: m = M().eval() exp_eq_scales = self.get_expected_eq_scales(m, x.detach().numpy()) prepared = prepare_fx( m, qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) prepared(x) convert_ref = _convert_equalization_ref(prepared) convert_ref(x) counter = 0 for node in convert_ref.graph.nodes: if 'equalization_scale' in node.name and node.op == 'get_attr': self.assertEqual(convert_ref.get_buffer(str(node.target)), exp_eq_scales[counter]) counter += 1
def test_input_weight_equalization_convert(self): """ Tests that the modified model for equalization (before quantization) returns the same output as the original model """ tests = [ SingleLayerLinearModel, LinearAddModel, TwoLayerLinearModel, SingleLayerFunctionalLinearModel, FunctionalLinearAddModel, TwoLayerFunctionalLinearModel ] x = torch.rand((5, 5)) for M in tests: m = M().eval() prepared = prepare_fx( copy.deepcopy(m), qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) output = prepared(x) convert_ref = _convert_equalization_ref(prepared) convert_ref_output = convert_ref(x) prepared = prepare_fx( m, qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) prepared(x) convert_fx(prepared) # Check if compile self.assertEqual(output, convert_ref_output)
def test_input_weight_equalization_activation_values(self): """ After applying the equalization functions check if the input observer's min/max values are as expected """ tests = [ SingleLayerLinearModel, TwoLayerLinearModel, SingleLayerFunctionalLinearModel ] x = torch.rand((5, 5)) torch.manual_seed(0) for M in tests: m = M().eval() exp_eq_scales = self.get_expected_eq_scales(m, x.detach().numpy()) exp_weights, exp_bias = self.get_expected_weights_bias( m, x.detach().numpy(), exp_eq_scales) exp_inp_act_vals = self.get_expected_inp_act_vals( m, x, exp_eq_scales, exp_weights, exp_bias) exp_weight_act_vals = self.get_expected_weight_act_vals( exp_weights) prepared = prepare_fx( m, qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) prepared(x) convert_ref = _convert_equalization_ref(prepared) convert_ref(x) modules = dict(convert_ref.named_modules(remove_duplicate=False)) inp_counter = 0 weight_counter = 0 for node in convert_ref.graph.nodes: if "weight" not in node.name and node.op == 'call_module' and \ isinstance(modules[str(node.target)], MinMaxObserver): # Check min/max values of input activation layers exp_min_val, exp_max_val = exp_inp_act_vals[inp_counter] self.assertEqual(modules[str(node.target)].min_val, exp_min_val) self.assertEqual(modules[str(node.target)].max_val, exp_max_val) inp_counter += 1 elif node.op == 'call_module' and isinstance( modules[str(node.target)], MinMaxObserver): # Check min/max values of weight activation layers assert ("weight" in node.name) exp_min_val, exp_max_val = exp_weight_act_vals[ weight_counter] self.assertEqual(modules[str(node.target)].min_val, exp_min_val) self.assertEqual(modules[str(node.target)].max_val, exp_max_val) weight_counter += 1
def test_input_weight_equalization_convert(self): """ Tests that the modified model for equalization (before quantization) returns the same output as the original model """ tests = [(SingleLayerLinearModel, 2), (LinearAddModel, 2), (TwoLayerLinearModel, 2), (SingleLayerFunctionalLinearModel, 2), (FunctionalLinearAddModel, 2), (TwoLayerFunctionalLinearModel, 2), (LinearReluModel, 2), (LinearReluLinearModel, 2), (LinearReluAddModel, 2), (FunctionalLinearReluModel, 2), (FunctionalLinearReluLinearModel, 2), (ConvModel, 4), (TwoLayerConvModel, 4), (SingleLayerFunctionalConvModel, 4), (TwoLayerFunctionalConvModel, 4), (ConvReluModel, 4), (ConvReluConvModel, 4), (ConvReluAddModel, 4), (FunctionalConvReluModel, 4), (FunctionalConvReluConvModel, 4)] for (M, ndim) in tests: m = M().eval() if ndim == 2: x = torch.rand((5, 5)) elif ndim == 4: x = torch.rand((16, 3, 224, 224)) prepared = prepare_fx( copy.deepcopy(m), qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) output = prepared(x) convert_ref = _convert_equalization_ref(prepared) convert_ref_output = convert_ref(x) prepared = prepare_fx( m, qconfig_dict, equalization_qconfig_dict=default_equalization_qconfig_dict) prepared(x) convert_fx(prepared) # Check if compile self.assertEqual(output, convert_ref_output)
def test_input_weight_equalization_convert(self): """ """ qconfig_dict = { "": None, "object_type": [(nn.Linear, default_qconfig), (nn.functional.linear, default_qconfig)] } default_equalization_qconfig_dict = { "": None, "object_type": [(nn.Linear, default_equalization_qconfig), (nn.functional.linear, default_equalization_qconfig)] } # Basic test with one linear layer class LinearModule(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(2, 2) def forward(self, x): return self.linear(x) # Test with two linear layer with a fp32 operation between class Linear2FP32Module(nn.Module): def __init__(self): super().__init__() self.linear1 = nn.Linear(2, 2) self.linear2 = nn.Linear(2, 2) def forward(self, x): x = self.linear1(x) x = torch.add(x, torch.tensor([1, 2])) x = self.linear2(x) return x tests = [(LinearModule, default_equalization_qconfig_dict), (Linear2FP32Module, default_equalization_qconfig_dict)] for (M, equalization_qconfig_dict) in tests: m = M().eval() x = torch.tensor([[1.0, 2.0], [2.0, 2.5], [4.5, 6.0]]) prepared = prepare_fx( m, qconfig_dict, equalization_qconfig_dict=equalization_qconfig_dict) output = prepared(x) convert_ref = _convert_equalization_ref(prepared) convert_ref_output = convert_ref(x) m = M().eval() prepared = prepare_fx( m, qconfig_dict, equalization_qconfig_dict=equalization_qconfig_dict) prepared(x) convert_fx(prepared) # Check if compile? self.assertEqual(output, convert_ref_output)