def test_acts_quant_params_linear(act1_type, act2_type, bn_out_stats): # prepare model: model = LinearBNSplitAct(act1_type, act2_type) stats = gen_stats_for_model(model) stats['bn']['output'] = bn_out_stats quantizer = PostTrainLinearQuantizer( model, model_activation_stats=deepcopy(stats)) quantizer.prepare_model(torch.randn(10, 10)) # get quant params: expected_quant_params_keys = { 'linear.output_zero_point', 'linear.output_scale', 'act1.output_zero_point', 'act1.output_scale', 'act2.output_zero_point', 'act2.output_scale' } assert set(quantizer.acts_quant_params) == expected_quant_params_keys quantizer.set_act_quant_param('linear.output_zero_point', 2.) quantizer.set_act_quant_param('linear.output_scale', 30.) assert model.linear.output_zero_point == 2. assert model.linear.output_scale == 30. expected_quant_param_linear_dict = { 'output_zero_point': torch.tensor(2.), 'output_scale': 30. } assert dict(model.linear.named_acts_quant_params() ) == expected_quant_param_linear_dict new_config = {'linear.output_zero_point': 4., 'act2.output_scale': 50} quantizer.update_acts_quant_params(new_config) assert model.linear.output_zero_point == 4 assert model.act2.output_scale == 50
def test_acts_quant_params_rnn(rnn_model): model = DummyWordLangModel(nn.Embedding(41, 20), rnn_model).cuda() stats = gen_stats_for_model(model) quantizer = PostTrainLinearQuantizer(model, model_activation_stats=deepcopy(stats)) dummy_input = torch.randint(0, 41, size=(79, 23)) quantizer.prepare_model(dummy_input) new_config = { 'rnn.rnn.cells.0.act_o.output_scale': 4, 'embedding.w_scale': torch.tensor(59.0) } quantizer.update_acts_quant_params(new_config) assert model.rnn.rnn.cells[0].act_o.output_scale == 4 assert model.embedding.w_scale == 59.0