def test_hawq_behaviour__if_method_returns_none(mocker, method_name,
expected_behavior):
config = HAWQConfigBuilder().with_sample_size([1, 1, 4,
4]).for_trial().build()
config['compression']['initializer']['range']['num_init_samples'] = 0
model = BasicConvTestModel()
mock_train_loader = mocker.stub()
mock_train_loader.batch_size = 1
device = 'cuda' if torch.cuda.is_available() else 'cpu'
config.register_extra_structs([
QuantizationPrecisionInitArgs(criterion_fn=mocker.stub(),
criterion=mocker.stub(),
data_loader=mock_train_loader,
device=device)
])
mocker.patch(
'nncf.common.initialization.batchnorm_adaptation.BatchnormAdaptationAlgorithm.run'
)
mocked_calc_traces = mocker.patch(
'nncf.torch.quantization.precision_init.hawq_init.HAWQPrecisionInitializer._calc_traces'
)
stub = mocker.stub()
stub.traces_order = TracesOrder([0])
mocked_calc_traces.return_value = stub
mocked_method = mocker.patch(
'nncf.torch.quantization.precision_init.hawq_init.HAWQPrecisionInitializer.'
+ method_name)
mocked_method.return_value = None
with expected_behavior:
create_compressed_model_and_algo_for_test(model, config)
def test_logarithm_scale_parameter(logarithm_scale_setting_1,
logarithm_scale_setting_2,
quantization_type):
for logarithm_scales in [[False, True], [True, False]]:
for symmetric in [False, True]:
model0, _ = create_compressed_model_and_algo_for_test(
TwoConvTestModel(),
get_config_for_logarithm_scale(
logarithm_scale=logarithm_scale_setting_1,
quantization_type=quantization_type))
model1, _ = create_compressed_model_and_algo_for_test(
TwoConvTestModel(),
get_config_for_logarithm_scale(
logarithm_scale=logarithm_scale_setting_2,
quantization_type=quantization_type))
sd0 = model0.state_dict()
model1.load_state_dict(sd0)
sd1 = model1.state_dict()
for k, v0 in sd0.items():
v1 = sd1[k] # pylint: disable=E1136
diff = (v1 - v0).abs().sum().item() / v1.numel()
assert diff < 1e-6, "symmetric {} logarithm_scales {} param {} is corrupted mean({}-{})={}".format(
symmetric, logarithm_scales, k, v0, v1, diff)
def test_can_restore_binary_mask_on_magnitude_algo_resume():
config = get_empty_config()
config['compression'] = {
"algorithm": "magnitude_sparsity",
"params": {
"weight_importance": "abs",
"schedule": "multistep",
"multistep_sparsity_levels": [0.3, 0.5]
}
}
sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
with torch.no_grad():
sparse_model(torch.ones([1, 1, 10, 10]))
config = get_empty_config()
config["compression"] = {"algorithm": "const_sparsity"}
const_sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
load_state(const_sparse_model, sparse_model.state_dict())
op = const_sparse_model.conv1.pre_ops['0']
PTTensorListComparator.check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.conv2.pre_ops['0']
PTTensorListComparator.check_equal(ref_mask_2, op.operand.binary_mask)
def test_load_state_interoperability(_algos, _model_wrapper, is_resume):
config_save = get_empty_config()
config_save['compression'] = [{
'algorithm': algo
} for algo in _algos['save_algos']]
register_bn_adaptation_init_args(config_save)
compressed_model_save, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config_save)
model_save = _model_wrapper['save_model'](compressed_model_save)
saved_model_state = model_save.state_dict()
ref_num_loaded = len(saved_model_state)
config_resume = get_empty_config()
config_resume['compression'] = [{
'algorithm': algo
} for algo in _algos['load_algos']]
register_bn_adaptation_init_args(config_resume)
compressed_model_resume, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config_resume)
model_resume = _model_wrapper['resume_model'](compressed_model_resume)
if not is_resume or (is_resume and _algos['is_resume_ok']):
act_num_loaded = load_state(model_resume, saved_model_state, is_resume)
if ('magnitude_sparsity' in _algos['load_algos'] or 'const_sparsity' in _algos['load_algos']) \
and 'rb_sparsity' in _algos['save_algos']:
# no need to load _mask and _uniform
ref_num_loaded -= 2
assert act_num_loaded == ref_num_loaded
else:
with pytest.raises(RuntimeError):
load_state(model_resume, saved_model_state, is_resume)
def test_load_state__with_resume_checkpoint(_resume_algos, _model_wrapper,
mocker):
config_save = get_empty_config()
config_save['compression'] = [{
'algorithm': algo
} for algo in _resume_algos['save_algos'] if algo != 'EMPTY']
register_bn_adaptation_init_args(config_save)
orig_model = BasicConvTestModel()
num_model_params = len(orig_model.state_dict())
model_save, compressed_ctrl_save = create_compressed_model_and_algo_for_test(
orig_model, config_save)
saved_model_state = model_save.state_dict()
saved_checkpoint = compressed_ctrl_save.get_compression_state()
ref_num_loaded = _resume_algos[
'ref_num_compression_params'] + num_model_params + 1 # padding_value
config_resume = get_empty_config()
config_resume['compression'] = [{
'algorithm': algo
} for algo in _resume_algos['load_algos'] if algo != 'EMPTY']
register_bn_adaptation_init_args(config_resume)
from nncf.torch.checkpoint_loading import KeyMatcher
key_matcher_run_spy = mocker.spy(KeyMatcher, 'run')
model, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(),
config_resume,
compression_state=saved_checkpoint)
load_state(model, saved_model_state, _resume_algos['is_strict'])
key_matcher_run_spy.assert_called_once()
act_num_loaded = len(key_matcher_run_spy.spy_return)
assert act_num_loaded == ref_num_loaded
def test_model_is_inited_with_own_device_by_default(
nncf_config_with_default_init_args, original_device):
if not torch.cuda.is_available() and 'cuda' in original_device:
pytest.skip("Skipping for CPU-only setups")
model = DeviceCheckingModel(original_device)
create_compressed_model_and_algo_for_test(
model, nncf_config_with_default_init_args)
def test_can_compress_with_config_and_resume_of_old_checkpoint():
model = SingleConv2dIdentityModel()
config = get_basic_quantization_config(
input_info={"sample_size": [1, 3, 100, 100]})
register_bn_adaptation_init_args(config)
create_compressed_model_and_algo_for_test(model,
config,
compression_state=old_style_sd)
def test_both_targets_assert():
config = get_basic_pruning_config()
config['compression']['algorithm'] = 'filter_pruning'
config['compression']['params']['pruning_target'] = 0.3
config['compression']['params']['pruning_flops_target'] = 0.5
model = PruningTestModel()
with pytest.raises(ValueError):
create_compressed_model_and_algo_for_test(model, config)
def test_context_independence(model_name, model_builder, input_size, _case_config):
config = get_basic_quantization_config(_case_config.quant_type, input_sample_sizes=input_size[0])
register_bn_adaptation_init_args(config)
compressed_models = [create_compressed_model_and_algo_for_test(model_builder[0](), config)[0],
create_compressed_model_and_algo_for_test(model_builder[1](), config)[0]]
for i, compressed_model in enumerate(compressed_models):
check_model_graph(compressed_model, model_name[i], _case_config.graph_dir)
def test_hawq_hw_vpu_config_e2e(_seed, dataset_dir, tmp_path):
config = HAWQConfigBuilder().for_vpu().liberal_mode().with_ratio(
1.5).build()
model = MobileNetV2(num_classes=10)
criterion = nn.CrossEntropyLoss()
if not dataset_dir:
dataset_dir = str(tmp_path)
train_loader, _ = create_test_dataloaders(config, dataset_dir)
config = register_default_init_args(config,
train_loader,
criterion=criterion)
create_compressed_model_and_algo_for_test(model, config)
def test_manual_single_conv(params):
config = params.nncf_config
register_bn_adaptation_init_args(config)
model = params.model
if params.expects_error:
with pytest.raises(ValueError):
create_compressed_model_and_algo_for_test(model, config)
else:
model, ctrl = create_compressed_model_and_algo_for_test(model, config)
path_to_dot = '{}.dot'.format(params.name)
graph_dir = os.path.join('quantized', 'hawq')
check_bitwidth_graph(ctrl, model, path_to_dot, graph_dir)
def test_is_overflow_fix_applied_model_resumed_correctly(tmp_path):
model = TwoConvTestModel()
nncf_config = get_config_for_export_mode(False)
compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, nncf_config)
compression_state = compression_ctrl.get_compression_state()
model_state_dict = compressed_model.state_dict()
# Must create new model as the previous one was somehow changed during create_compressed_model_and_algo_for_test()
model = TwoConvTestModel()
compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, nncf_config, compression_state=compression_state)
load_state(compressed_model, model_state_dict, is_resume=True)
are_symmetric_fq_nodes_are_exported_correct_with_overflow_fix(
tmp_path, compression_ctrl)
def test_frozen_layers(_nncf_caplog, params):
model = params.create_frozen_model()
config = params.create_config()
register_bn_adaptation_init_args(config)
if params.raising_error:
with pytest.raises(RuntimeError):
__, _ = create_compressed_model_and_algo_for_test(model, config)
else:
__, _ = create_compressed_model_and_algo_for_test(model, config)
are_frozen_layers_mentioned = 'Frozen layers' in _nncf_caplog.text
if params.printing_warning:
assert are_frozen_layers_mentioned
else:
assert not are_frozen_layers_mentioned
def test_evolution_env_default_params():
model = PruningTestModel()
config = create_default_legr_config()
train_loader = create_ones_mock_dataloader(config)
val_loader = create_ones_mock_dataloader(config)
train_steps_fn = lambda *x: None
validate_fn = lambda *x: (0, 0)
nncf_config = register_default_init_args(config,
train_loader=train_loader,
train_steps_fn=train_steps_fn,
val_loader=val_loader,
validate_fn=validate_fn)
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
evolution_env = compression_ctrl.legr.env
assert evolution_env.loss_as_reward is True
assert evolution_env.prune_target == 0.5
assert evolution_env.steps == 200
assert evolution_env.train_loader == train_loader
assert evolution_env.val_loader == val_loader
assert evolution_env.train_fn == train_steps_fn
assert evolution_env.validate_fn == validate_fn
assert evolution_env.config == nncf_config
def test_evolution_env_setting_params():
steps_ref = 100
prune_target_ref = 0.1
train_optimizer = partial(optim.Adam)
model = PruningTestModel()
config = create_default_legr_config()
config['compression']['params']['legr_params'] = {}
config['compression']['params']['legr_params']['train_steps'] = steps_ref
config['compression']['params']['legr_params'][
'max_pruning'] = prune_target_ref
train_loader = create_ones_mock_dataloader(config)
val_loader = create_ones_mock_dataloader(config)
train_steps_fn = lambda *x: None
validate_fn = lambda *x: (0, 0)
nncf_config = register_default_init_args(
config,
train_loader=train_loader,
train_steps_fn=train_steps_fn,
val_loader=val_loader,
validate_fn=validate_fn,
legr_train_optimizer=train_optimizer)
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, nncf_config)
evolution_env = compression_ctrl.legr.env
assert evolution_env.prune_target == prune_target_ref
assert evolution_env.steps == steps_ref
assert evolution_env.train_optimizer == train_optimizer
def create_finetuned_lenet_model_and_dataloader(config,
eval_fn,
finetuning_steps,
learning_rate=1e-3):
with set_torch_seed():
train_loader = create_ones_mock_dataloader(config, num_samples=10)
model = LeNet()
for param in model.parameters():
nn.init.uniform_(param, a=0.0, b=0.01)
data_loader = iter(train_loader)
optimizer = SGD(model.parameters(), lr=learning_rate)
for _ in range(finetuning_steps):
optimizer.zero_grad()
x, y_gt = next(data_loader)
y = model(x)
loss = F.mse_loss(y.sum(), y_gt)
loss.backward()
optimizer.step()
config = register_default_init_args(
config,
train_loader=train_loader,
model_eval_fn=partial(eval_fn, train_loader=train_loader))
model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
return model, train_loader, compression_ctrl
def run_actual(model: nn.Module,
config: NNCFConfig,
inference_type: str,
mock_dataloader: Iterable,
ngpus_per_node=None) -> Tuple[List[torch.Tensor], NNCFNetwork]:
config = get_kd_config(config)
model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
if inference_type == 'DDP':
model = post_compression_test_distr_init(compression_ctrl, config,
ngpus_per_node, model)
elif inference_type in ('DP', 'single_GPU'):
if inference_type == 'DP':
model = torch.nn.DataParallel(model)
optimizer = SGD(model.parameters(), lr=1e-02, weight_decay=1e-02)
model.train()
output_storage = []
for _, (input_, __) in enumerate(mock_dataloader):
input_ = input_.to(next(model.parameters()).device)
output = model(input_)
output_storage.append(output)
loss = compression_ctrl.loss()
optimizer.zero_grad()
loss.backward()
optimizer.step()
return output_storage, model
def test_knowledge_distillation_outputs_containers_parsing():
mse = torch.nn.MSELoss()
input_size = [1, 1, 8, 8]
model = ContainersOutputsModel(input_size)
fill_params_of_model_by_normal(model)
dumped_orig_model = deepcopy(model)
sparsity_level = 0.3
batch_size = 1 if torch.cuda.device_count(
) == 0 else torch.cuda.device_count()
config = get_kd_config(
get_sparsity_config_with_sparsity_init(
get_basic_magnitude_sparsity_config(input_sample_size=input_size),
sparsity_level))
model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
model.train()
mock_dataloader = create_ones_mock_dataloader(
config, num_samples=torch.cuda.device_count(), batch_size=batch_size)
compression_ctrl.scheduler.epoch_step()
for _, (input_, __) in enumerate(mock_dataloader):
input_ = input_.to(next(model.parameters()).device)
outputs = model(input_)
kd_outputs = dumped_orig_model(input_)
reference_kd_loss = mse(outputs['xa'], kd_outputs['xa']) + \
mse(outputs['xb_and_xc'][0], kd_outputs['xb_and_xc'][0]) + \
mse(outputs['xb_and_xc'][1], kd_outputs['xb_and_xc'][1])
actual_kd_loss = compression_ctrl.loss()
assert torch.allclose(reference_kd_loss, actual_kd_loss)
def test_valid_masks_for_bn_after_concat(prune_bn):
config = get_basic_pruning_config(input_sample_size=[1, 1, 8, 8])
config['compression']['algorithm'] = 'filter_pruning'
config['compression']['params']['prune_batch_norms'] = prune_bn
config['compression']['params']['prune_first_conv'] = True
config['compression']['pruning_init'] = 0.5
model = PruningTestModelConcatBN()
pruned_model, _ = create_compressed_model_and_algo_for_test(model, config)
bn_modules = [pruned_model.bn, pruned_model.bn1, pruned_model.bn2]
for bn_module in bn_modules:
if prune_bn:
# Check that mask was applied for batch_norm module
mask = bn_module.pre_ops['0'].op.binary_filter_pruning_mask
assert sum(mask) == len(mask) * 0.5
else:
# Check that no mask was added to the layer
assert len(bn_module.pre_ops) == 0
# Check output mask of concat layers
ref_concat_masks = [[0] * 8 + [1] * 8 + [0] * 8 + [1] * 8,
[1] * 8 + [0] * 16 + [1] * 8 + [0] * 8 + [1] * 8]
graph = pruned_model.get_original_graph()
for i, node in enumerate(graph.get_nodes_by_types(['cat'])):
assert np.allclose(node.data['output_mask'].tensor.numpy(),
ref_concat_masks[i])
def test_magnitude_algo_set_independently_sparsity_level_for_one_module():
module_name_conv1 = 'MagnitudeTestModel/NNCFConv2d[conv1]/conv2d_0'
module_name_conv2 = 'MagnitudeTestModel/NNCFConv2d[conv2]/conv2d_0'
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = {"sparsity_level_setting_mode": 'local'}
sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(MagnitudeTestModel(), config)
sparse_info_conv1 = [sparse_info for sparse_info in compression_ctrl.sparsified_module_info\
if sparse_info.module_node_name == module_name_conv1]
sparse_info_conv2 = [sparse_info for sparse_info in compression_ctrl.sparsified_module_info\
if sparse_info.module_node_name == module_name_conv2]
compression_ctrl.set_sparsity_level(0.5, sparse_info_conv1[0])
weights_conv1 = sparse_model.conv1.weight
weights_conv2 = sparse_model.conv2.weight
count_nonzero_conv1 = sparse_model.conv1.pre_ops['0'].operand.apply_binary_mask(weights_conv1).nonzero().size(0)
count_param_conv1 = weights_conv1.view(-1).size(0)
assert count_param_conv1 - count_nonzero_conv1 == 4 # 8 * 0.5
compression_ctrl.set_sparsity_level(0.3, sparse_info_conv2[0])
count_nonzero_conv1 = sparse_model.conv1.pre_ops['0'].operand.apply_binary_mask(weights_conv1).nonzero().size(0)
count_param_conv1 = weights_conv1.view(-1).size(0)
count_nonzero_conv2 = sparse_model.conv2.pre_ops['0'].operand.apply_binary_mask(weights_conv2).nonzero().size(0)
count_param_conv2 = weights_conv2.view(-1).size(0)
assert count_param_conv1 - count_nonzero_conv1 == 4 # 8 * 0.5
assert count_param_conv2 - count_nonzero_conv2 == 6 # ~ 18 * 0.3
def test_ad_hoc_range_init_does_not_replace_parameter_tensors(
self, wrap_dataloader, quant_type):
config = create_config()
config["compression"].update({
"activations": {
"mode": quant_type
},
"weights": {
"mode": quant_type
}
})
data_loader = self.create_dataloader(wrap_dataloader, config)
config.register_extra_structs([QuantizationRangeInitArgs(data_loader)])
model = TwoConvTestModel()
quant_model, quant_ctrl = create_compressed_model_and_algo_for_test(
model, config)
param_name_vs_id = {
name: id(tnsr)
for name, tnsr in quant_model.named_parameters()
}
quant_ctrl.init_range()
for name, param in quant_model.named_parameters():
assert param_name_vs_id[name] == id(param)
def test_can_create_magnitude_sparse_algo__with_defaults():
model = MagnitudeTestModel()
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = \
{'schedule': 'multistep'}
sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(deepcopy(model), config)
assert isinstance(compression_ctrl, MagnitudeSparsityController)
assert compression_ctrl.scheduler.current_sparsity_level == approx(0.1)
assert len(list(sparse_model.modules())) == 12
_, sparse_model_conv = check_correct_nncf_modules_replacement(model, sparse_model)
i = 0
nncf_stats = compression_ctrl.statistics()
for layer_info in nncf_stats.magnitude_sparsity.thresholds:
assert layer_info.threshold == approx(0.24, 0.1)
# pylint: disable=protected-access
assert isinstance(compression_ctrl._weight_importance_fn, type(normed_magnitude))
for sparse_module in sparse_model_conv.values():
store = []
ref_mask = torch.ones_like(sparse_module.weight) if i == 0 else ref_mask_2
i += 1
for op in sparse_module.pre_ops.values():
if isinstance(op, UpdateWeight) and isinstance(op.operand, BinaryMask):
assert torch.allclose(op.operand.binary_mask, ref_mask)
assert op.__class__.__name__ not in store
store.append(op.__class__.__name__)
def test_is_pytorch_output_the_same_as_onnx_qdq_overflow_fix_applied(
tmp_path, model):
nncf_config = get_config_for_export_mode(True)
nncf_config.update({"input_info": {"sample_size": [1, 1, 20, 20]}})
compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, nncf_config)
onnx_checkpoint_path = str(tmp_path / 'model.onnx')
compression_ctrl.export_model(onnx_checkpoint_path)
input_tensors = [
np.random.normal(size=[1, 1, 20, 20]),
np.random.uniform(size=[1, 1, 20, 20]),
100 * np.random.normal(size=[1, 1, 20, 20]),
100 * np.random.uniform(size=[1, 1, 20, 20])
]
for input_tensor in input_tensors:
torch_input = torch.tensor(input_tensor, dtype=torch.float32)
with torch.no_grad():
torch_out = compressed_model(torch_input)
# ONNXRuntime
sess = rt.InferenceSession(onnx_checkpoint_path)
input_name = sess.get_inputs()[0].name
onnx_out = sess.run(None,
{input_name: input_tensor.astype(np.float32)})[0]
assert np.allclose(torch_out.numpy(), onnx_out, rtol=1e-5, atol=1e-3)
def test_can_quantize_inputs_for_sparsity_plus_quantization():
model = BasicConvTestModel()
config = get_basic_sparsity_plus_quantization_config()
register_bn_adaptation_init_args(config)
sparse_quantized_model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
assert isinstance(compression_ctrl,
CompositeCompressionAlgorithmController)
sparse_quantized_model_conv = get_all_modules_by_type(
sparse_quantized_model, 'NNCFConv2d')
nncf_module = next(iter(sparse_quantized_model_conv.values()))
assert len(
nncf_module.pre_ops) == 2 # 1x weight sparsifier + 1x weight quantizer
assert isinstance(nncf_module.pre_ops['0'], UpdateWeight)
assert isinstance(nncf_module.pre_ops['0'].op, RBSparsifyingWeight)
assert isinstance(nncf_module.pre_ops['1'], UpdateWeight)
assert isinstance(nncf_module.pre_ops['1'].op, SymmetricQuantizer)
input_quantizer = get_all_modules(sparse_quantized_model)[
f'NNCFNetwork/ModuleDict[{EXTERNAL_QUANTIZERS_STORAGE_NAME}]']
assert len(input_quantizer) == 1
assert isinstance(list(input_quantizer.values())[0], SymmetricQuantizer)
def test_staged_scheduler_with_empty_quantization():
config = get_squeezenet_quantization_config()
config['compression'].update({
'params': {
"activations_quant_start_epoch": 1,
"weights_quant_start_epoch": 2,
}
})
register_bn_adaptation_init_args(config)
model = squeezenet1_1(num_classes=10, dropout=0)
model, algo = create_compressed_model_and_algo_for_test(model, config)
scheduler = algo.scheduler
for module in algo.all_quantizations.values():
assert not module.is_enabled_quantization()
scheduler.epoch_step()
for module in algo.all_quantizations.values():
assert not module.is_enabled_quantization()
scheduler.epoch_step()
for wq_info in algo.weight_quantizers.values():
assert not wq_info.quantizer_module_ref.is_enabled_quantization()
for aq_info in algo.non_weight_quantizers.values():
assert aq_info.quantizer_module_ref.is_enabled_quantization()
scheduler.epoch_step()
for module in algo.all_quantizations.values():
assert module.is_enabled_quantization()
def test_export_stacked_bi_lstm(tmp_path):
p = LSTMTestSizes(3, 3, 3, 3)
config = get_empty_config(
input_sample_sizes=[1, p.hidden_size, p.input_size])
config['compression'] = {'algorithm': 'quantization'}
register_bn_adaptation_init_args(config)
# TODO: batch_first=True fails with building graph: ambiguous call to mul or sigmoid
test_rnn = NNCF_RNN('LSTM',
input_size=p.input_size,
hidden_size=p.hidden_size,
num_layers=2,
bidirectional=True,
batch_first=False)
model, algo = create_compressed_model_and_algo_for_test(test_rnn, config)
test_path = str(tmp_path.joinpath('test.onnx'))
algo.export_model(test_path)
assert os.path.exists(test_path)
onnx_num = 0
model = onnx.load(test_path)
# pylint: disable=no-member
for node in model.graph.node:
if node.op_type == 'FakeQuantize':
onnx_num += 1
assert onnx_num == 54
def test_can_do_sparsity_freeze_epoch():
def compare_binary_mask(ref_sparse_module_info, sparse_module_info):
for ref_sparse_layer, sparse_layer in zip(ref_sparse_module_info, sparse_module_info):
if (ref_sparse_layer.operand.binary_mask != sparse_layer.operand.binary_mask).view(-1).sum() != 0:
return False
return True
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {"algorithm": "magnitude_sparsity",
"sparsity_init": 0.1,
"params": {"sparsity_target": 0.9,
"sparsity_target_epoch": 3,
"sparsity_freeze_epoch": 3}}
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
sparsified_minfo_before_update = deepcopy(compression_ctrl.sparsified_module_info)
compression_ctrl.scheduler.epoch_step() # update binary_masks
compression_ctrl.scheduler.epoch_step() # update binary_masks
compression_ctrl.scheduler.epoch_step() # update binary_masks, freeze binary_masks
sparsified_minfo_after_update = deepcopy(compression_ctrl.sparsified_module_info)
assert not compare_binary_mask(sparsified_minfo_after_update, sparsified_minfo_before_update)
compression_ctrl.scheduler.epoch_step() # don't update binary_masks
sparsified_minfo_after_freeze = deepcopy(compression_ctrl.sparsified_module_info)
assert compare_binary_mask(sparsified_minfo_after_update, sparsified_minfo_after_freeze)
def test_loss_outputs_parsing():
mse = torch.nn.MSELoss()
input_size = [1, 1, 8, 8]
model = PartlyNonDifferentialOutputsModel(input_size)
fill_params_of_model_by_normal(model)
dumped_orig_model = deepcopy(model)
sparsity_level = 0.3
batch_size = 1 if torch.cuda.device_count(
) == 0 else torch.cuda.device_count()
config = get_kd_config(
get_sparsity_config_with_sparsity_init(
get_basic_magnitude_sparsity_config(input_sample_size=input_size),
sparsity_level))
model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
model.train()
mock_dataloader = create_ones_mock_dataloader(
config, num_samples=torch.cuda.device_count(), batch_size=batch_size)
compression_ctrl.scheduler.epoch_step()
for _, (input_, __) in enumerate(mock_dataloader):
input_ = input_.to(next(model.parameters()).device)
outputs = model(input_)
kd_outputs = dumped_orig_model(input_)
loss_outputs = []
for tensor1, tensor2 in zip(outputs, kd_outputs):
if tensor1.requires_grad:
loss_outputs.append((tensor1, tensor2))
reference_kd_loss = sum(
[mse(item[0], item[1]) for item in loss_outputs])
actual_kd_loss = compression_ctrl.loss()
assert torch.allclose(reference_kd_loss, actual_kd_loss)
def test_can_resume_with_algo_mixing(mocker, is_strict):
desc = TestPrecisionInitDesc().config_with_all_inits()
all_quantization_init_spies = desc.setup_init_spies(mocker)
sparsity_config = get_basic_sparsity_config()
sparsity_config['target_device'] = 'TRIAL'
config = desc.config
quantization_section = config['compression']
config['compression'] = [{
'algorithm': 'const_sparsity'
}, quantization_section]
_, compression_ctrl = create_compressed_model_and_algo_for_test(
desc.model_creator(), sparsity_config)
compression_state = compression_ctrl.get_compression_state()
config = register_default_init_args(
config, train_loader=create_ones_mock_dataloader(config))
fn = partial(create_compressed_model_and_algo_for_test,
desc.model_creator(),
config,
compression_state=compression_state)
if is_strict:
with pytest.raises(RuntimeError):
fn()
else:
_, compression_ctrl = fn()
for m in all_quantization_init_spies:
m.assert_called()
desc.check_precision_init(compression_ctrl.child_ctrls[1])
def test_check_default_algo_params():
"""
Test for default algorithm params. Creating empty config and check for valid default
parameters.
"""
# Creating algorithm with empty config
config = get_basic_pruning_config()
config['compression']['algorithm'] = 'filter_pruning'
model = PruningTestModel()
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
assert isinstance(compression_ctrl, FilterPruningController)
scheduler = compression_ctrl.scheduler
# Check default algo params
assert compression_ctrl.prune_first is False
assert compression_ctrl.prune_batch_norms is True
assert compression_ctrl.prune_downsample_convs is False
assert compression_ctrl.filter_importance is l2_filter_norm
assert compression_ctrl.ranking_type == 'unweighted_ranking'
assert compression_ctrl.pruning_quota == 0.9
assert compression_ctrl.all_weights is False
# Check default scheduler params
assert isinstance(scheduler, ExponentialPruningScheduler)
