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)
