| # Owner(s): ["oncall: quantization"] |
| |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torch.ao.nn.intrinsic.quantized as nniq |
| import torch.ao.nn.quantized as nnq |
| from torch.ao.quantization import default_qconfig |
| from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver |
| from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx |
| from torch.ao.quantization.fx._equalize import ( |
| _InputEqualizationObserver, |
| _WeightEqualizationObserver, |
| calculate_equalization_scale, |
| default_equalization_qconfig, |
| _convert_equalization_ref, |
| get_layer_sqnr_dict, |
| get_equalization_qconfig_dict, |
| ) |
| |
| from torch.testing._internal.common_quantization import ( |
| NodeSpec as ns, |
| QuantizationTestCase, |
| SingleLayerLinearModel, |
| TwoLayerLinearModel, |
| LinearAddModel, |
| SingleLayerFunctionalLinearModel, |
| TwoLayerFunctionalLinearModel, |
| FunctionalLinearAddModel, |
| ConvModel, |
| TwoLayerConvModel, |
| SingleLayerFunctionalConvModel, |
| TwoLayerFunctionalConvModel, |
| skipIfNoFBGEMM, |
| LinearReluModel, |
| LinearReluLinearModel, |
| LinearReluAddModel, |
| FunctionalLinearReluModel, |
| FunctionalLinearReluLinearModel, |
| ConvReluModel, |
| ConvReluConvModel, |
| ConvReluAddModel, |
| FunctionalConvReluModel, |
| FunctionalConvReluConvModel, |
| ) |
| |
| # Standard Libraries |
| import copy |
| import numpy as np |
| |
| # Testing utils |
| from hypothesis import given |
| from hypothesis import strategies as st |
| |
| |
| default_qconfig_dict = {"": default_qconfig} |
| |
| specific_qconfig_dict = { |
| "": None, |
| "object_type": [(nn.Linear, default_qconfig), |
| (F.linear, default_qconfig), |
| (nn.ReLU, default_qconfig), |
| (F.relu, default_qconfig), |
| (nn.Conv2d, default_qconfig), |
| (F.conv2d, default_qconfig)] |
| } |
| |
| default_equalization_qconfig_dict = { |
| "": None, |
| "object_type": [(nn.Linear, default_equalization_qconfig), |
| (F.linear, default_equalization_qconfig), |
| (nn.ReLU, default_equalization_qconfig), |
| (F.relu, default_equalization_qconfig), |
| (nn.Conv2d, default_equalization_qconfig), |
| (F.conv2d, default_equalization_qconfig)] |
| } |
| |
| |
| class TestEqualizeFx(QuantizationTestCase): |
| def channel_minmax(self, input, axis=1): |
| ''' Finds the min/max of inputs associated with a specific channel |
| ''' |
| size_of_tensor_dim = input.ndim |
| axis_list = list(range(size_of_tensor_dim)) |
| axis_list.remove(axis) |
| axis_list.sort(reverse=True) |
| |
| mins = input.copy() |
| maxs = input.copy() |
| for a in axis_list: |
| mins = mins.min(a) |
| maxs = maxs.max(a) |
| |
| return (mins, maxs) |
| |
| @given(ndim=st.sampled_from((2, 3, 4, 5)), |
| input_qdtype=st.sampled_from((torch.qint8, torch.quint8)), |
| input_qscheme=st.sampled_from((torch.per_tensor_affine, torch.per_tensor_symmetric)), |
| weight_qdtype=st.sampled_from((torch.qint8, torch.quint8)), |
| weight_qscheme=st.sampled_from((torch.per_channel_affine, torch.per_channel_symmetric, |
| torch.per_channel_affine_float_qparams))) |
| def test_input_weight_eq_observer(self, ndim, input_qdtype, input_qscheme, weight_qdtype, weight_qscheme): |
| sizes = [] |
| for _ in range((ndim - 1) * 2): |
| sizes.append(np.random.randint(2, 10)) |
| |
| channel = np.random.randint(1, 10) |
| if ndim == 2: |
| x = np.random.random(size=(sizes[0], channel)) |
| w = np.random.random(size=(sizes[1], channel)) |
| elif ndim == 3: |
| x = np.random.random(size=(sizes[0], channel, sizes[1])) |
| w = np.random.random(size=(sizes[2], channel, sizes[3])) |
| elif ndim == 4: |
| x = np.random.random(size=(sizes[0], channel, sizes[1], sizes[2])) |
| w = np.random.random(size=(sizes[3], channel, sizes[4], sizes[5])) |
| elif ndim == 5: |
| x = np.random.random(size=(sizes[0], channel, sizes[1], sizes[2], sizes[3])) |
| w = np.random.random(size=(sizes[4], channel, sizes[5], sizes[6], sizes[7])) |
| |
| x = (x * 10).round(decimals=2).astype(np.float32) |
| w = (w * 10).round(decimals=2).astype(np.float32) |
| |
| input_eq_obs = _InputEqualizationObserver(dtype=input_qdtype, qscheme=input_qscheme) |
| weight_eq_obs = _WeightEqualizationObserver(dtype=weight_qdtype, qscheme=weight_qscheme) |
| |
| ret_x = input_eq_obs(torch.tensor(x)) |
| ret_w = weight_eq_obs(torch.tensor(w)) |
| self.assertEqual((ret_x, ret_w), (x, w)) |
| |
| # Check the min/max input columns are correct |
| ref_min_inputs, ref_max_inputs = self.channel_minmax(x) |
| min_inputs, max_inputs = input_eq_obs.get_input_minmax() |
| self.assertEqual(min_inputs, torch.tensor(ref_min_inputs, dtype=torch.float32)) |
| self.assertEqual(max_inputs, torch.tensor(ref_max_inputs, dtype=torch.float32)) |
| |
| # Check the min/max weight columns are correct |
| ref_min_weights_col, ref_max_weights_col = self.channel_minmax(w) |
| min_weights_col, max_weights_col = weight_eq_obs.get_weight_col_minmax() |
| self.assertEqual(min_weights_col, torch.tensor(ref_min_weights_col, dtype=torch.float32)) |
| self.assertEqual(max_weights_col, torch.tensor(ref_max_weights_col, dtype=torch.float32)) |
| |
| # Check the equalization scale is correct |
| equalization_scale = calculate_equalization_scale(input_eq_obs, weight_eq_obs) |
| ref_equalization_scale = np.sqrt((ref_max_weights_col - ref_min_weights_col) / |
| (ref_max_inputs - ref_min_inputs)) |
| self.assertEqual(equalization_scale, torch.tensor(ref_equalization_scale, dtype=torch.float32)) |
| |
| input_eq_obs.set_equalization_scale(equalization_scale) |
| weight_eq_obs.set_equalization_scale(equalization_scale) |
| |
| # Check the input scale/zero-point values |
| min_input_scaled, max_input_scaled = input_eq_obs.calculate_scaled_minmax() |
| input_quant_obs = MinMaxObserver(dtype=input_qdtype, qscheme=input_qscheme) |
| input_quant_obs.min_val = min_input_scaled |
| input_quant_obs.max_val = max_input_scaled |
| input_qparams = input_quant_obs.calculate_qparams() |
| |
| ref_min_input_scaled = np.min(ref_min_inputs * ref_equalization_scale) |
| ref_min_input_scaled = min(0, ref_min_input_scaled) |
| ref_max_input_scaled = np.max(ref_max_inputs * ref_equalization_scale) |
| ref_max_input_scaled = max(0, ref_max_input_scaled) |
| |
| if input_qscheme == torch.per_tensor_symmetric: |
| ref_scale = 2 * max(abs(ref_min_input_scaled), ref_max_input_scaled) / 255 |
| ref_zero_point = 0 if input_qdtype is torch.qint8 else 128 |
| else: |
| ref_scale = (ref_max_input_scaled - ref_min_input_scaled) / 255 |
| quant_min = -128 if input_qdtype is torch.qint8 else 0 |
| quant_max = 127 if input_qdtype is torch.qint8 else 255 |
| ref_zero_point = quant_min - np.round(ref_min_input_scaled / ref_scale) |
| np.clip(ref_zero_point, quant_min, quant_max) |
| |
| self.assertEqual(input_qparams[0].item(), ref_scale, atol=1e-5, rtol=0) |
| self.assertEqual(input_qparams[1].item(), ref_zero_point) |
| |
| # During input-weight equalization, we will scale the weights so that |
| # the following weight quantized observer will have the correct scaled qparams |
| # Check the weight scale/zero-point values of the quantized observer |
| weight_quant_obs = PerChannelMinMaxObserver(ch_axis=1, dtype=weight_qdtype, qscheme=weight_qscheme) |
| |
| # Scale the weights for input-weight equalization |
| new_shape = [1] * w.ndim |
| new_shape[1] = w.shape[1] |
| ref_w_scaled = w * np.reciprocal(ref_equalization_scale.reshape(tuple(new_shape))) |
| |
| w = torch.tensor(w) |
| new_shape[1] = w.size(1) |
| w_scaled = torch.mul(w, torch.reciprocal(equalization_scale.view(new_shape))) |
| |
| self.assertEqual(w_scaled, ref_w_scaled) |
| |
| # Call forward on the weight quantization observer |
| weight_quant_obs(w_scaled) |
| |
| # Check the min/max weight rows are correct |
| ref_min_weights_scaled, ref_max_weights_scaled = self.channel_minmax(ref_w_scaled) |
| self.assertEqual(weight_quant_obs.min_val, torch.tensor(ref_min_weights_scaled, dtype=torch.float32)) |
| self.assertEqual(weight_quant_obs.max_val, torch.tensor(ref_max_weights_scaled, dtype=torch.float32)) |
| |
| weight_qparams = weight_quant_obs.calculate_qparams() |
| |
| if weight_qscheme == torch.per_channel_symmetric: |
| ref_min_weights_scaled = np.minimum(np.zeros(ref_min_weights_scaled.shape), ref_min_weights_scaled) |
| ref_max_weights_scaled = np.maximum(np.zeros(ref_max_weights_scaled.shape), ref_max_weights_scaled) |
| |
| ref_scales = 2 * np.maximum(np.abs(ref_min_weights_scaled), ref_max_weights_scaled) / 255 |
| ref_zero_points = np.zeros_like( |
| ref_scales) if weight_qdtype is torch.qint8 else np.ones_like(ref_scales) * 128 |
| elif weight_qscheme == torch.per_channel_affine_float_qparams: |
| ref_scales = (ref_max_weights_scaled - ref_min_weights_scaled) / 255 |
| ref_scales = np.where(ref_scales > 1e-7, ref_scales, np.ones_like(ref_scales)) |
| ref_zero_points = -1 * ref_min_weights_scaled / ref_scales |
| else: |
| ref_min_weights_scaled = np.minimum(np.zeros_like(ref_min_weights_scaled), ref_min_weights_scaled) |
| ref_max_weights_scaled = np.maximum(np.zeros_like(ref_max_weights_scaled), ref_max_weights_scaled) |
| |
| ref_scales = (ref_max_weights_scaled - ref_min_weights_scaled) / 255 |
| ref_zero_points = -128 if weight_qdtype is torch.qint8 else 0 |
| ref_zero_points = ref_zero_points - np.round(ref_min_weights_scaled / ref_scales) |
| |
| self.assertEqual(weight_qparams[0], torch.tensor( |
| ref_scales, dtype=weight_qparams[0].dtype), rtol=1e-5, atol=0.0001) |
| self.assertEqual(weight_qparams[1], torch.tensor( |
| ref_zero_points, dtype=weight_qparams[1].dtype), rtol=1e-5, atol=1) |
| |
| def test_input_weight_equalization_prepare(self): |
| """ Tests that graphs created after prepare_fx is as expected |
| """ |
| |
| single_nn_layer_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 1, |
| ns.call_module(MinMaxObserver): 2, |
| } |
| |
| two_nn_layer_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 2, |
| ns.call_module(MinMaxObserver): 3, |
| } |
| |
| single_F_layer_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 1, |
| ns.call_module(_WeightEqualizationObserver): 1, |
| ns.call_module(MinMaxObserver): 3, |
| } |
| |
| two_F_layer_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 2, |
| ns.call_module(_WeightEqualizationObserver): 2, |
| ns.call_module(MinMaxObserver): 5, |
| } |
| |
| fp_F_layer_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 2, |
| ns.call_module(_WeightEqualizationObserver): 2, |
| ns.call_module(MinMaxObserver): 6, |
| } |
| |
| tests = [(SingleLayerLinearModel, single_nn_layer_node_occurrence), |
| (TwoLayerLinearModel, two_nn_layer_node_occurrence), |
| (TwoLayerFunctionalLinearModel, two_F_layer_node_occurrence), |
| (FunctionalLinearAddModel, fp_F_layer_node_occurrence), |
| (LinearReluModel, single_nn_layer_node_occurrence), |
| (LinearReluLinearModel, two_nn_layer_node_occurrence), |
| (FunctionalLinearReluModel, single_F_layer_node_occurrence), |
| (FunctionalLinearReluLinearModel, two_F_layer_node_occurrence), |
| (ConvModel, single_nn_layer_node_occurrence), |
| (TwoLayerConvModel, two_nn_layer_node_occurrence), |
| (TwoLayerFunctionalConvModel, two_F_layer_node_occurrence), |
| (ConvReluModel, single_nn_layer_node_occurrence), |
| (ConvReluConvModel, two_nn_layer_node_occurrence), |
| (FunctionalConvReluModel, single_F_layer_node_occurrence), |
| (FunctionalConvReluConvModel, two_F_layer_node_occurrence)] |
| |
| for (M, node_occurrence) in tests: |
| m = M().eval() |
| example_inputs = m.get_example_inputs() |
| prepared = prepare_fx( |
| m, |
| specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict) |
| self.checkGraphModuleNodes(prepared, expected_node_occurrence=node_occurrence) |
| |
| def test_input_weight_equalization_branching(self): |
| """ Tests that graphs containing branches are prepared correctly. |
| Specifically, equalization observers should not be inserted in front of |
| branches in which both initial layers in the branches plan to be |
| quantized. |
| """ |
| |
| # Tests that we do not add an equalization observer due to both initial |
| # nodes in the branch containing layers that need to be equalized. |
| # Note that this should print out 2 warning messages for not being able |
| # to equalize layers linear1 and linear1 because it is part of a branch |
| class TestBranchingWithoutEqualizationModel(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.linear1 = nn.Linear(5, 5) |
| self.linear2 = nn.Linear(5, 5) |
| |
| def forward(self, x): |
| y = self.linear1(x) |
| z = self.linear2(x) |
| return torch.add(y, z) |
| |
| no_eq_branching_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 0, |
| ns.call_module(MinMaxObserver): 3, |
| } |
| |
| m = TestBranchingWithoutEqualizationModel().eval() |
| example_inputs = (torch.rand(1, 5),) |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict) |
| self.checkGraphModuleNodes(prepared, expected_node_occurrence=no_eq_branching_node_occurrence) |
| |
| # Tests that we will add an equalization observer because there is only |
| # one initial node in the branch that needs to be equalized |
| class TestBranchingWithEqualizationModel(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.linear1 = nn.Linear(5, 5) |
| |
| def forward(self, x): |
| y = self.linear1(x) |
| z = torch.add(x, 5) |
| return torch.add(y, z) |
| |
| eq_branching_node_occurrence = { |
| ns.call_module(_InputEqualizationObserver): 1, |
| ns.call_module(MinMaxObserver): 2, |
| } |
| |
| m = TestBranchingWithEqualizationModel().eval() |
| example_inputs = (torch.randn(1, 5),) |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict) |
| self.checkGraphModuleNodes(prepared, expected_node_occurrence=eq_branching_node_occurrence) |
| |
| @skipIfNoFBGEMM |
| 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)) |
| |
| example_inputs = (x,) |
| prepared = prepare_fx( |
| copy.deepcopy(m), |
| specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict |
| ) |
| output = prepared(x) |
| |
| convert_ref = _convert_equalization_ref(prepared) |
| convert_ref_output = convert_ref(x) |
| |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict) |
| prepared(x) |
| convert_fx(prepared) # Check if compile |
| self.assertEqual(output, convert_ref_output) |
| |
| def calculate_equalization_scale_ref(self, x, w): |
| """ Calculates the equalization scale based on the input and weight |
| """ |
| min_inputs = x.min(axis=0) |
| max_inputs = x.max(axis=0) |
| |
| min_weights_col = w.min(axis=0) |
| max_weights_col = w.max(axis=0) |
| |
| equalization_scale = np.sqrt((max_weights_col - min_weights_col) / |
| (max_inputs - min_inputs)) |
| return equalization_scale |
| |
| def get_expected_eq_scales(self, model, x): |
| """ For each module in the graph, we want to calculate the equalization |
| scale at that point. This only works for models containing single or |
| connected linear layers. |
| """ |
| exp_eq_scales = [] |
| for _, module in model.named_children(): |
| weight = module.weight.detach().numpy() |
| bias = module.bias.detach().numpy() |
| |
| eq_scale = self.calculate_equalization_scale_ref(x, weight) |
| exp_eq_scales.append(eq_scale) |
| |
| x = x @ weight.T + bias |
| |
| return exp_eq_scales |
| |
| 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()) |
| |
| example_inputs = (x,) |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict) |
| prepared(*example_inputs) |
| 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)).reshape(-1), exp_eq_scales[counter]) |
| counter += 1 |
| |
| def get_expected_weights_bias(self, model, x, exp_eq_scales): |
| """ For each module in the graph, we want to calculate the expected |
| scaled weight and bias values. This only works for models containing |
| single or connected linear layers. |
| """ |
| exp_weights = [] |
| exp_bias = [] |
| for i, (_, module) in enumerate(model.named_children()): |
| weight = module.weight.detach().numpy() |
| bias = module.bias.detach().numpy() |
| |
| scaled_weight = weight * np.reciprocal(exp_eq_scales[i]) |
| scaled_bias = bias |
| if i + 1 < len(exp_eq_scales): |
| scaled_weight = (scaled_weight.T * exp_eq_scales[i + 1]).T |
| scaled_bias = (scaled_bias.T * exp_eq_scales[i + 1]).T |
| |
| exp_weights.append(scaled_weight) |
| exp_bias.append(scaled_bias) |
| |
| x = x @ weight.T + bias |
| |
| return exp_weights, exp_bias |
| |
| 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) |
| |
| example_inputs = (x,) |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=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 get_expected_inp_act_vals(self, model, x, exp_eq_scales, exp_weights, exp_bias): |
| """ For each module in the graph, we want to calculate the expected |
| min/max values for every input activation node. This only works for |
| models containing only single or connected linear layers. |
| """ |
| x = x * exp_eq_scales[0] |
| |
| exp_inp_activation_vals = [] |
| for i, _ in enumerate(model.named_children()): |
| exp_inp_activation_vals.append((x.min(), x.max())) |
| x = x @ exp_weights[i].T + exp_bias[i] |
| |
| exp_inp_activation_vals.append((x.min(), x.max())) |
| return exp_inp_activation_vals |
| |
| def get_expected_weight_act_vals(self, exp_weights): |
| """ For each module in the graph, we want to calculate the expected |
| min/max values for every weight activation node. This is assuming that |
| the weight observers are all MinMaxObservers. |
| """ |
| |
| exp_weight_activation_vals = [] |
| for w in exp_weights: |
| exp_weight_activation_vals.append((w.min(), w.max())) |
| |
| return exp_weight_activation_vals |
| |
| 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) |
| |
| example_inputs = (x,) |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=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: |
| users = list(node.users) |
| if node.op == 'call_module' and isinstance(modules[str(node.target)], MinMaxObserver): |
| if len(users) == 1 and users[0].target == torch.nn.functional.linear and users[0].args[1] == node: |
| # Check min/max values of weight activation layers |
| 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 |
| else: |
| # 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 |
| |
| |
| def check_orig_and_eq_graphs(self, orig_model, eq_model): |
| """ Given a non-equalized model and an equalized model, check that the |
| graphs are structured in the same way, except the equalized model has |
| additional 'equalization_scale' and 'mul' nodes. |
| """ |
| orig_idx = 0 |
| orig_nodes = list(orig_model.graph.nodes) |
| orig_modules = dict(orig_model.named_modules(remove_duplicate=False)) |
| |
| eq_idx = 0 |
| eq_nodes = list(eq_model.graph.nodes) |
| eq_modules = dict(eq_model.named_modules(remove_duplicate=False)) |
| |
| while orig_idx < len(orig_nodes) and eq_idx < len(eq_nodes): |
| if 'equalization_scale' in eq_nodes[eq_idx].name and 'mul' in eq_nodes[eq_idx + 1].name: |
| # Skip the equalization and mul nodes |
| eq_idx += 2 |
| continue |
| elif orig_nodes[orig_idx].op != eq_nodes[eq_idx].op: |
| return False |
| elif orig_nodes[orig_idx].op == 'call_module': |
| # Check that the type of call_modules are the same (ex. nn.Linear, MinMaxObserver) |
| orig_node = orig_nodes[orig_idx] |
| eq_node = eq_nodes[eq_idx] |
| if type(orig_modules[orig_node.target]) is not type(eq_modules[eq_node.target]): |
| return False |
| elif orig_nodes[orig_idx].op == 'call_function': |
| # Check that the call_functions are the same (ex. F.linear) |
| orig_node = orig_nodes[orig_idx] |
| eq_node = eq_nodes[eq_idx] |
| if orig_node.target != eq_node.target: |
| return False |
| |
| eq_idx += 1 |
| orig_idx += 1 |
| |
| return True |
| |
| @skipIfNoFBGEMM |
| def test_input_weight_equalization_graphs(self): |
| """ Tests that the modified model for equalization has the same graph |
| structure as the model without equalization (before and after |
| quantization). |
| """ |
| |
| linear_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize') |
| ] |
| |
| linearAdd_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize'), |
| ns.call_function(torch.add), |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize') |
| ] |
| |
| linear2_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Linear), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalLinear_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.linear), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalLinearAdd_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.linear), |
| ns.call_method('dequantize'), |
| ns.call_function(torch.add), |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.linear), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalLinear2_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.linear), |
| ns.call_function(torch.ops.quantized.linear), |
| ns.call_method('dequantize') |
| ] |
| |
| linearRelu_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nniq.LinearReLU), |
| ns.call_method('dequantize') |
| ] |
| |
| linearReluLinear_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nniq.LinearReLU), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalLinearRelu_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.linear_relu), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalLinearReluLinear_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.linear_relu), |
| ns.call_function(torch.ops.quantized.linear), |
| ns.call_method('dequantize') |
| ] |
| |
| conv_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Conv2d), |
| ns.call_method('dequantize') |
| ] |
| |
| conv2_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Conv2d), |
| ns.call_module(nnq.Conv2d), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalConv_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.conv2d), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalConv2_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.conv2d), |
| ns.call_function(torch.ops.quantized.conv2d), |
| ns.call_method('dequantize') |
| ] |
| |
| convRelu_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nniq.ConvReLU2d), |
| ns.call_method('dequantize') |
| ] |
| |
| convReluConv_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nniq.ConvReLU2d), |
| ns.call_module(nnq.Conv2d), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalConvRelu_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.conv2d_relu), |
| ns.call_method('dequantize') |
| ] |
| |
| functionalConvReluConv_node_list = [ |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_function(torch.ops.quantized.conv2d_relu), |
| ns.call_function(torch.ops.quantized.conv2d), |
| ns.call_method('dequantize') |
| ] |
| |
| tests = [(SingleLayerLinearModel, linear_node_list), |
| (LinearAddModel, linearAdd_node_list), |
| (TwoLayerLinearModel, linear2_node_list), |
| (SingleLayerFunctionalLinearModel, functionalLinear_node_list), |
| (FunctionalLinearAddModel, functionalLinearAdd_node_list), |
| (TwoLayerFunctionalLinearModel, functionalLinear2_node_list), |
| (LinearReluModel, linearRelu_node_list), |
| (LinearReluLinearModel, linearReluLinear_node_list), |
| (FunctionalLinearReluModel, functionalLinearRelu_node_list), |
| (FunctionalLinearReluLinearModel, functionalLinearReluLinear_node_list), |
| (ConvModel, conv_node_list), |
| (TwoLayerConvModel, conv2_node_list), |
| (SingleLayerFunctionalConvModel, functionalConv_node_list), |
| (TwoLayerFunctionalConvModel, functionalConv2_node_list), |
| (ConvReluModel, convRelu_node_list), |
| (ConvReluConvModel, convReluConv_node_list), |
| (FunctionalConvReluModel, functionalConvRelu_node_list), |
| (FunctionalConvReluConvModel, functionalConvReluConv_node_list)] |
| |
| for (M, node_list) in tests: |
| m = M().eval() |
| example_inputs = m.get_example_inputs() |
| prepared = prepare_fx( |
| m, specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict) |
| equalized_quantized_model = convert_fx(prepared) |
| |
| # Check the order of nodes in the graph |
| self.checkGraphModuleNodes(equalized_quantized_model, expected_node_list=node_list) |
| |
| @skipIfNoFBGEMM |
| def test_input_weight_equalization_results(self): |
| """ Tests that for small models, the results of quantized models that |
| have been equalized are very close to models that have not been equalized. |
| """ |
| |
| tests = [SingleLayerLinearModel, TwoLayerLinearModel, LinearAddModel, |
| SingleLayerFunctionalLinearModel, TwoLayerFunctionalLinearModel] |
| |
| x = torch.rand((5, 5)) |
| for M in tests: |
| m = M().eval() |
| |
| # No equalization |
| example_inputs = (x,) |
| prepared = prepare_fx( |
| copy.deepcopy(m), |
| specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config={}) |
| prepared(x) |
| quantized = convert_fx(prepared) # Check if compile |
| quantized_output = quantized(x) |
| |
| # With equalization |
| prepared = prepare_fx( |
| copy.deepcopy(m), |
| specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=default_equalization_qconfig_dict |
| ) |
| prepared(x) |
| equalized_and_quantized = convert_fx(prepared) # Check if compile |
| equalized_and_quantized_output = equalized_and_quantized(x) |
| self.assertEqual(quantized_output, equalized_and_quantized_output, rtol=1e-5, atol=0.1) |
| |
| @skipIfNoFBGEMM |
| def test_selective_equalization(self): |
| """ Tests that we are able to run numeric suite on the equalized model |
| and construct a valid equalization_config equalizing only the top |
| 4 layers with the highest quantization errors. |
| """ |
| |
| torch.manual_seed(1) |
| |
| class M(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.bot = torch.nn.Sequential(torch.nn.Linear(5, 5)) |
| self.top = torch.nn.Sequential(torch.nn.Linear(5, 5)) |
| |
| def forward(self, x): |
| x = self.bot(x) |
| x = torch.add(x, 5) |
| x = self.top(x) |
| return x |
| |
| float_model = M().eval() |
| # Hard coded so that the top layer has a higher quantization error |
| x = torch.tensor([[0.0642, 0.7824, 0.4255, 0.7106, 0.5957], |
| [0.8373, 0.8851, 0.8229, 0.0212, 0.8987], |
| [0.9077, 0.7538, 0.4530, 0.5772, 0.1376], |
| [0.0690, 0.9002, 0.7998, 0.2768, 0.8985], |
| [0.0282, 0.5068, 0.6725, 0.1829, 0.5480]]) |
| |
| # Quantize the float model |
| example_inputs = (x,) |
| prepared_model = prepare_fx( |
| copy.deepcopy(float_model), |
| specific_qconfig_dict, |
| example_inputs=example_inputs |
| ) |
| prepared_model(x) |
| quantized_model = convert_fx(copy.deepcopy(prepared_model)) |
| |
| # Get the SQNR between the float and quantized model |
| layer_to_sqnr_dict = get_layer_sqnr_dict(copy.deepcopy(prepared_model), quantized_model, x) |
| |
| # Construct the equalization_qconfig_dict equalizing layers with the highest |
| # quantization errors |
| selective_equalization_qconfig_dict = get_equalization_qconfig_dict(layer_to_sqnr_dict, 1) |
| |
| # Create the selectively equalized model |
| prepared_model = prepare_fx( |
| copy.deepcopy(float_model), |
| specific_qconfig_dict, |
| example_inputs=example_inputs, |
| _equalization_config=selective_equalization_qconfig_dict, |
| ) |
| prepared_model(x) |
| equalized_model = convert_fx(prepared_model) |
| |
| node_list = [ |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize'), |
| ns.call_function(torch.add), |
| ns.call_function(torch.mul), |
| ns.call_function(torch.quantize_per_tensor), |
| ns.call_module(nnq.Linear), |
| ns.call_method('dequantize') |
| ] |
| |
| # Check the order of nodes in the graph |
| self.checkGraphModuleNodes(equalized_model, expected_node_list=node_list) |
