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
