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_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_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_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 get_multistep_normed_abs_config(): config = get_basic_magnitude_sparsity_config() compression_config = config['compression'] compression_config['params'] = { 'schedule': 'multistep', 'weight_importance': 'normed_abs', 'multistep_steps': [1, 3], 'multistep_sparsity_levels': [0.1, 0.5, 0.9] } return config
def test_magnitude_sparse_algo_sets_threshold(weight_importance, sparsity_level, threshold): model = MagnitudeTestModel() config = get_basic_magnitude_sparsity_config() config['compression']['params'] = {'schedule': 'multistep', 'weight_importance': weight_importance} _, compression_ctrl = create_compressed_model_and_algo_for_test(model, config) if sparsity_level: compression_ctrl.set_sparsity_level(sparsity_level) nncf_stats = compression_ctrl.statistics() for layer_info in nncf_stats.magnitude_sparsity.thresholds: assert layer_info.threshold == pytest.approx(threshold, 0.01)
def test_magnitude_algo_set_binary_mask_on_forward(): config = get_basic_magnitude_sparsity_config() config['compression']['params'] = {'weight_importance': 'abs'} sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(MagnitudeTestModel(), config) compression_ctrl.set_sparsity_level(0.3) with torch.no_grad(): sparse_model(torch.ones([1, 1, 10, 10])) op = sparse_model.conv1.pre_ops['0'] PTTensorListComparator.check_equal(ref_mask_1, op.operand.binary_mask) op = sparse_model.conv2.pre_ops['0'] PTTensorListComparator.check_equal(ref_mask_2, op.operand.binary_mask)
def test_knowledge_distillation_loss_types(kd_loss_type: str, scale, temperature): torch.manual_seed(2) if kd_loss_type == 'softmax': def kd_loss_fn(ref_outputs, compressed_model_outputs) -> torch.Tensor: return scale * -(nn.functional.log_softmax( compressed_model_outputs / temperature, dim=1) * nn.functional.softmax( ref_outputs / temperature, dim=1)).mean() * ( temperature * temperature * compressed_model_outputs.shape[1]) else: def mse_loss(x, y): return scale * F.mse_loss(x, y) kd_loss_fn = mse_loss input_size = [1, 100] batch_size = 1 if torch.cuda.device_count( ) == 0 else torch.cuda.device_count() model = nn.Sequential( nn.Linear(in_features=input_size[-1], out_features=10), nn.Sigmoid()) fill_params_of_model_by_normal(model) dumped_orig_model = deepcopy(model) sparsity_level = 0.5 config = get_kd_config(get_sparsity_config_with_sparsity_init( get_basic_magnitude_sparsity_config(input_sample_size=input_size), sparsity_level), kd_type=kd_loss_type, scale=scale, temperature=temperature) config['compression'][-1]['type'] = kd_loss_type 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 = kd_loss_fn(kd_outputs, outputs) actual_kd_loss = compression_ctrl.loss() assert torch.allclose(reference_kd_loss, actual_kd_loss)
def test_runner(num_steps, learning_rate, reference_metric): runner = PTAccuracyAwareTrainingRunner(accuracy_aware_training_params={}, dump_checkpoints=False) input_sample_size = [1, 1, LeNet.INPUT_SIZE[-1], LeNet.INPUT_SIZE[-1]] config = get_basic_magnitude_sparsity_config( input_sample_size=input_sample_size) model, train_loader, compression_ctrl = create_initialized_lenet_model_and_dataloader( config) def train_fn(compression_ctrl, model, epoch, optimizer, lr_scheduler, train_loader=train_loader): with set_torch_seed(): train_loader = iter(train_loader) for _ in range(num_steps): compression_ctrl.scheduler.step() optimizer.zero_grad() x, y_gt = next(train_loader) y = model(x) loss = F.mse_loss(y.sum(), y_gt) loss.backward() optimizer.step() def validate_fn(model, epoch, train_loader=train_loader): with set_torch_seed(): train_loader = iter(train_loader) loss = 0 with torch.no_grad(): for _ in range(num_steps): x, y_gt = next(train_loader) y = model(x) loss += F.mse_loss(y.sum(), y_gt) return loss.item() def configure_optimizers_fn(): optimizer = SGD(model.parameters(), lr=learning_rate) return optimizer, None runner.initialize_training_loop_fns(train_fn, validate_fn, configure_optimizers_fn, None) runner.reset_training() runner.train_epoch(model, compression_ctrl) metric_value = runner.validate(model) assert metric_value == pytest.approx(reference_metric, 1e-3)
def test_magnitude_algo_can_calculate_sparsity_rate_for_one_sparsified_module(): module_node_name_conv1 = 'MagnitudeTestModel/NNCFConv2d[conv1]/conv2d_0' config = get_basic_magnitude_sparsity_config() config['compression']['params'] = {"sparsity_level_setting_mode": 'local'} _, 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_node_name_conv1] compression_ctrl.set_sparsity_level(0.5, sparse_info_conv1[0]) nncf_stats = compression_ctrl.statistics() sparse_model_stats = nncf_stats.magnitude_sparsity.model_statistics module_name_to_sparsity_level_map = { s.name: s.sparsity_level for s in sparse_model_stats.sparsified_layers_summary } module_name = sparse_info_conv1[0].module_node_name assert pytest.approx(module_name_to_sparsity_level_map[module_name], 1e-2) == 0.5
def test_knowledge_distillation_training_process(inference_type: str): if not torch.cuda.is_available() and not inference_type == 'cpu': pytest.skip("Skipping CUDA test cases for CPU only setups") torch.manual_seed(1) input_size = [1, 1, 8, 8] sparsity_level = 0.3 config = get_sparsity_config_with_sparsity_init( get_basic_magnitude_sparsity_config(input_sample_size=input_size), sparsity_level) if inference_type == 'DDP': ngpus_per_node = torch.cuda.device_count() config.world_size = ngpus_per_node torch.multiprocessing.spawn(run_test_training, nprocs=ngpus_per_node, args=(config, inference_type, ngpus_per_node), join=True) else: run_test_training(None, config, inference_type, None)
def test_magnitude_algo_can_calculate_correct_stats_for_local_mode(): module_node_name_conv1 = 'MagnitudeTestModel/NNCFConv2d[conv1]/conv2d_0' module_node_name_conv2 = 'MagnitudeTestModel/NNCFConv2d[conv2]/conv2d_0' config = get_basic_magnitude_sparsity_config() config['compression']['params'] = {"sparsity_level_setting_mode": 'local'} _, 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_node_name_conv1] sparse_info_conv2 = [sparse_info for sparse_info in compression_ctrl.sparsified_module_info\ if sparse_info.module_node_name == module_node_name_conv2] compression_ctrl.set_sparsity_level(0.5, sparse_info_conv1[0]) compression_ctrl.set_sparsity_level(0.3, sparse_info_conv2[0]) nncf_stats = compression_ctrl.statistics() expected = [(module_node_name_conv1, 0.3344823), (module_node_name_conv2, 0.2191864)] for idx, layer_info in enumerate(nncf_stats.magnitude_sparsity.thresholds): expected_name, expected_threshold = expected[idx] assert layer_info.name == expected_name assert pytest.approx(layer_info.threshold) == expected_threshold
def test_kd_incompatible_output_shapes_handling(outputs_dim_numbers_list, kd_type, is_zero): """ Checks ignorance behavior (kd loss is zero) for different model output shape sizes :param dim_numbers_list: a list of dim numbers of model output tensors examples: dim_numbers_list = [4] -> model outputs should be [*, *, *, *] dim_numbers_list = [4, 2] -> model outputs should be ([*, *, *, *], [*, *]) :param kd_type: type of knowledge distillation loss "param is_zero: if given type of model outputs should be ignored than kd loss value should be zero """ input_size = [1, 2, 3, 4] config = get_kd_config(get_sparsity_config_with_sparsity_init( get_basic_magnitude_sparsity_config(input_sample_size=input_size), sparsity_init=0.5), kd_type=kd_type) model = CustomOutputWeightedModel(input_size, outputs_dim_numbers_list) compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) compression_ctrl.scheduler.epoch_step() compressed_model.train() input_ = torch.normal(0, std=0.5, size=input_size) compressed_model.forward(input_) kd_loss = compression_ctrl.loss() assert torch.allclose(kd_loss, torch.zeros([1])) == is_zero
def test_can_not_create_magnitude_algo__without_steps(): config = get_basic_magnitude_sparsity_config() config['compression']['params'] = {'schedule': 'multistep', 'multistep_sparsity_levels': [0.1]} with pytest.raises(ValueError): _, _ = create_compressed_model_and_algo_for_test(MockModel(), config)
def test_can_create_magnitude_algo__without_levels(): config = get_basic_magnitude_sparsity_config() config['compression']['params'] = {'schedule': 'multistep', 'multistep_steps': [1]} _, compression_ctrl = create_compressed_model_and_algo_for_test(MockModel(), config) assert compression_ctrl.scheduler.current_sparsity_level == approx(0.1)
def test_can_set_compression_rate_for_magnitude_sparse_algo(): config = get_basic_magnitude_sparsity_config() _, compression_ctrl = create_compressed_model_and_algo_for_test(MagnitudeTestModel(), config) compression_ctrl.compression_rate = 0.65 assert pytest.approx(compression_ctrl.compression_rate, 1e-2) == 0.65
def test_adaptive_compression_training_loop(max_accuracy_degradation, final_compression_rate, reference_final_metric, num_steps=10, learning_rate=1e-3, initial_training_phase_epochs=5, patience_epochs=3, init_finetuning_steps=10): def validate_fn(model, epoch=0, train_loader=None): with set_torch_seed(): train_loader = iter(train_loader) loss = torch.FloatTensor([0]) with torch.no_grad(): for _ in range(num_steps): x, y_gt = next(train_loader) y = model(x) loss += F.mse_loss(y.sum(), y_gt) return 1 - loss.item() input_sample_size = [1, 1, LeNet.INPUT_SIZE[-1], LeNet.INPUT_SIZE[-1]] config = get_basic_magnitude_sparsity_config( input_sample_size=input_sample_size) params = { "initial_training_phase_epochs": initial_training_phase_epochs, "patience_epochs": patience_epochs, } params.update(max_accuracy_degradation) accuracy_aware_config = { "accuracy_aware_training": { "mode": "adaptive_compression_level", "params": params } } config.update(accuracy_aware_config) model, train_loader, compression_ctrl = create_finetuned_lenet_model_and_dataloader( config, validate_fn, init_finetuning_steps) def train_fn(compression_ctrl, model, optimizer, train_loader=train_loader, **kwargs): with set_torch_seed(): train_loader = iter(train_loader) for _ in range(num_steps): compression_ctrl.scheduler.step() optimizer.zero_grad() x, y_gt = next(train_loader) y = model(x) loss = F.mse_loss(y.sum(), y_gt) loss.backward() optimizer.step() def configure_optimizers_fn(): optimizer = SGD(model.parameters(), lr=learning_rate) return optimizer, None acc_aware_training_loop = AdaptiveCompressionTrainingLoop( config, compression_ctrl) model = acc_aware_training_loop.run( model, train_epoch_fn=train_fn, validate_fn=partial(validate_fn, train_loader=train_loader), configure_optimizers_fn=configure_optimizers_fn) assert compression_ctrl.compression_rate == pytest.approx( final_compression_rate, 1e-3) assert validate_fn(model, train_loader=train_loader) == pytest.approx( reference_final_metric, 1e-4)
def test_can_not_set_sparsity_more_than_one_for_magnitude_sparse_algo(): config = get_basic_magnitude_sparsity_config() _, compression_ctrl = create_compressed_model_and_algo_for_test(MagnitudeTestModel(), config) with pytest.raises(AttributeError): compression_ctrl.set_sparsity_level(1) compression_ctrl.set_sparsity_level(1.2)