def test_load_state_interoperability(_algos, _models, is_resume):
config_save = get_empty_config()
config_save['compression'] = [{
'algorithm': algo,
'params': {}
} for algo in _algos['save_algos']]
algo_save = create_test_compression_algo(config_save,
_models['save_model'])
model_save = algo_save.model
saved_model_state = model_save.state_dict()
ref_num_loaded = len(saved_model_state)
config_resume = get_empty_config()
config_resume['compression'] = [{
'algorithm': algo,
'params': {}
} for algo in _algos['load_algos']]
algo_resume = create_test_compression_algo(config_resume,
_models['resume_model'])
model_resume = algo_resume.model
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_can_restore_binary_mask_on_magnitude_algo_resume():
config = get_empty_config()
config['compression'] = {
"algorithm": "magnitude_sparsity",
"weight_importance": "abs",
"params": {
"schedule": "multistep",
"sparsity_levels": [0.3, 0.5]
}
}
magnitude_algo = create_compression_algorithm(MagnitudeTestModel(), config)
sparse_model = magnitude_algo.model
with torch.no_grad():
sparse_model(torch.ones([1, 1, 10, 10]))
config = get_empty_config()
config["compression"] = {"algorithm": "const_sparsity"}
const_algo = create_compression_algorithm(MagnitudeTestModel(), config)
const_sparse_model = const_algo.model
load_state(const_sparse_model, sparse_model.state_dict())
op = const_sparse_model.conv1.pre_ops['0']
check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.conv2.pre_ops['0']
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']]
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']]
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_can_restore_binary_mask_on_magnitude_quant_algo_resume():
config = get_empty_config()
config["compression"] = [
{"algorithm": "magnitude_sparsity", "weight_importance": "abs",
"params": {"schedule": "multistep", "sparsity_levels": [0.3, 0.5]}},
{"algorithm": "quantization"}]
reset_context('orig')
reset_context('quantized_graphs')
magnitude_quant_algo = create_compression_algorithm(MagnitudeTestModel(), config)
# load_state doesn't support CPU + Quantization
sparse_model = torch.nn.DataParallel(magnitude_quant_algo.model)
sparse_model.cuda()
with torch.no_grad():
sparse_model(torch.ones([1, 1, 10, 10]))
reset_context('orig')
reset_context('quantized_graphs')
config = get_empty_config()
config["compression"] = [{"algorithm": "const_sparsity"}, {"algorithm": "quantization"}]
const_algo = create_compression_algorithm(MagnitudeTestModel(), config)
const_sparse_model = const_algo.model
load_state(const_sparse_model, sparse_model.state_dict())
op = const_sparse_model.module.conv1.pre_ops['0']
check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.module.conv2.pre_ops['0']
check_equal(ref_mask_2, op.operand.binary_mask)
def test_export_stacked_bi_lstm(tmp_path):
p = LSTMTestSizes(3, 3, 3, 3)
config = get_empty_config(input_sample_size=(1, p.hidden_size,
p.input_size))
config['compression'] = {'algorithm': 'quantization'}
# 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 == 50
def test_can_choose_scheduler(algo, schedule_type, scheduler_class):
config = get_empty_config()
config['compression']['algorithm'] = algo
config['compression']["params"]["schedule"] = schedule_type
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, scheduler_class)
def test_can_create_rb_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression']['algorithm'] = 'rb_sparsity'
config['compression']["params"]["schedule"] = 'adaptive'
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, AdaptiveSparsityScheduler)
def test_can_create_const_sparse_algo__with_default():
model = BasicConvTestModel()
config = get_empty_config()
config["compression"] = {"algorithm": "const_sparsity"}
compression_algo = create_compression_algorithm(deepcopy(model), config)
assert isinstance(compression_algo, ConstSparsity)
sparse_model = compression_algo.model
assert len(list(sparse_model.modules())) == 6
model_conv = get_all_modules_by_type(model, 'Conv2d')
sparse_model_conv = get_all_modules_by_type(sparse_model, 'NNCFConv2d')
assert len(model_conv) == len(sparse_model_conv)
for module_name in model_conv:
scope = module_name.split('/')
scope[-1] = scope[-1].replace('Conv2d', 'NNCFConv2d')
sparse_module_name = '/'.join(scope)
assert sparse_module_name in sparse_model_conv
store = []
sparse_module = sparse_model_conv[sparse_module_name]
for op in sparse_module.pre_ops.values():
if isinstance(op, UpdateWeight) and isinstance(
op.operand, BinaryMask):
ref_mask = torch.ones_like(sparse_module.weight)
assert torch.allclose(op.operand.binary_mask, ref_mask)
assert op.__class__.__name__ not in store
store.append(op.__class__.__name__)
def test_can_create_rb_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression']['algorithm'] = 'rb_sparsity'
config['compression']["params"]["schedule"] = 'adaptive'
compression_algo = create_compression_algorithm(BasicConvTestModel(),
config)
assert isinstance(compression_algo.scheduler, AdaptiveSparsityScheduler)
def test_can_choose_scheduler(algo, schedule_type, scheduler_class):
config = get_empty_config()
config['compression']['algorithm'] = algo
config['compression']["params"]["schedule"] = schedule_type
compression_algo = create_compression_algorithm(BasicConvTestModel(),
config)
assert isinstance(compression_algo.scheduler, scheduler_class)
def test_magnitude_algo_binary_masks_are_applied():
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {'algorithm': "magnitude_sparsity"}
compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
minfo_list = compression_ctrl.sparsified_module_info # type: List[SparseModuleInfo]
minfo = minfo_list[0] # type: SparseModuleInfo
minfo.operand.binary_mask = torch.ones_like(minfo.module.weight) # 1x1x2x2
input_ = torch.ones(size=(1, 1, 5, 5))
ref_output_1 = -4 * torch.ones(size=(2, 4, 4))
output_1 = compressed_model(input_)
assert torch.all(torch.eq(output_1, ref_output_1))
minfo.operand.binary_mask[0][0][0][1] = 0
minfo.operand.binary_mask[1][0][1][0] = 0
ref_output_2 = -3 * torch.ones_like(ref_output_1)
output_2 = compressed_model(input_)
assert torch.all(torch.eq(output_2, ref_output_2))
minfo.operand.binary_mask[1][0][0][1] = 0
ref_output_3 = ref_output_2.clone()
ref_output_3[1] = -2 * torch.ones_like(ref_output_1[1])
output_3 = compressed_model(input_)
assert torch.all(torch.eq(output_3, ref_output_3))
def get_basic_quantization_config_with_hw_config_type(hw_config_type,
input_sample_size):
config = get_empty_config(input_sample_sizes=input_sample_size)
config["hw_config_type"] = hw_config_type
config["compression"] = {
"algorithm": "quantization",
}
return config
def create_config():
config = get_empty_config()
config['compression'] = {
'algorithm': 'quantization',
'initializer': {
'num_init_steps': 1
}
}
return config
def get_basic_quantization_config(quantization_type, input_sample_size):
config = get_empty_config(input_sample_size=input_sample_size)
config["compression"] = {"algorithm": "quantization",
"activations": {
"mode": quantization_type
},
"weights": {
"mode": quantization_type
}}
return config
def test_can_create_sparse_scheduler__with_defaults(self, algo):
config = get_empty_config()
config['compression']['algorithm'] = algo
config['compression']["params"]["schedule"] = 'polynomial'
compression_algo = create_compression_algorithm(BasicConvTestModel(), config)
scheduler = compression_algo.scheduler
assert scheduler.initial_sparsity == 0
assert scheduler.max_sparsity == 0.5
assert scheduler.max_step == 90
assert scheduler.sparsity_training_steps == 100
def test_can_not_create_magnitude_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression'] = {
'algorithm': 'magnitude_sparsity',
"params": {
"schedule": 'adaptive'
}
}
with pytest.raises(TypeError):
_, _ = create_compressed_model_and_algo_for_test(MockModel(), config)
def test_number_of_calling_fq_for_lstm(self):
p = LSTMTestSizes(1, 1, 1, 5)
num_layers = 2
bidirectional = True
num_directions = 2 if bidirectional else 1
bias = True
batch_first = False
config = get_empty_config(input_sample_size=(p.seq_length, p.batch,
p.input_size))
config['compression'] = {
'algorithm': 'quantization',
'quantize_inputs': True
}
test_data = TestLSTMCell.generate_lstm_data(p,
num_layers,
num_directions,
bias=bias,
batch_first=batch_first)
test_rnn = NNCF_RNN('LSTM',
input_size=p.input_size,
hidden_size=p.hidden_size,
num_layers=num_layers,
bidirectional=bidirectional,
bias=bias,
batch_first=batch_first)
TestLSTM.set_ref_lstm_weights(test_data, test_rnn, num_layers,
num_directions, bias)
test_hidden = TestLSTM.get_test_lstm_hidden(test_data)
model, algo = create_compressed_model_and_algo_for_test(
test_rnn, config)
class Counter:
def __init__(self):
self.count = 0
def next(self):
self.count += 1
def hook(model, input_, counter):
counter.next()
counters = {}
for name, quantizer in algo.all_quantizations.items():
counter = Counter()
counters[name] = counter
quantizer.register_forward_pre_hook(partial(hook, counter=counter))
_ = model(test_data.x, test_hidden)
assert model.get_graph().get_nodes_count(
) == 107 # NB: may always fail in debug due to superfluous 'cat' nodes
assert len(counters) == 50
for counter in counters.values():
assert counter.count == p.seq_length
def test_can_create_sparse_scheduler__with_defaults(self, algo):
config = get_empty_config()
config['compression']['algorithm'] = algo
config['compression']["params"]["schedule"] = 'polynomial'
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
scheduler = compression_ctrl.scheduler
assert scheduler.initial_sparsity == 0
assert scheduler.max_sparsity == 0.5
assert scheduler.max_step == 90
assert scheduler.sparsity_training_steps == 100
def test_sparse_network(self, model_name, model_builder, input_size, algo, params):
model = model_builder()
from nncf.layers import NNCF_MODULES_MAP
sparsifiable_modules = list(NNCF_MODULES_MAP.values())
ref_num_sparsed = len(get_all_modules_by_type(model, sparsifiable_modules))
config = get_empty_config(input_sample_size=input_size)
config["compression"] = {"algorithm": algo, "params": params}
compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
assert ref_num_sparsed == len(compression_ctrl.sparsified_module_info)
check_model_graph(compressed_model, model_name, algo)
def test_scale_and_sign_init_for_quant_algo():
model = TwoConvTestModel()
config = get_empty_config()
config['compression'] = {
'algorithm': 'quantization',
'initializer': {
'num_init_steps': 1
}
}
reset_context('orig')
reset_context('quantized_graphs')
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
input_sample_size = config.input_sample_size
class OnesDatasetMock:
def __init__(self, input_size):
self.input_size = input_size
super().__init__()
def __getitem__(self, index):
return torch.ones(self.input_size), torch.ones(1)
def __len__(self):
return 1
data_loader = torch.utils.data.DataLoader(OnesDatasetMock(
input_sample_size[1:]),
batch_size=1,
num_workers=1,
shuffle=False)
compression_algo.initialize(data_loader)
model_conv = get_all_modules_by_type(model, 'Quantize')
ref_table = {
'.*Sequential\\[0\\].*UpdateWeight.*': (True, 1),
'.*Sequential\\[1\\].*UpdateWeight. *': (False, 1),
'.*activation_quantizers.*Sequential\\[0\\].*': (True, 4),
'.*activation_quantizers.*Sequential\\[1\\].*': (True, 24)
}
for name, module in model_conv.items():
for pattern, ref_values in ref_table.items():
match = re.search(pattern, name)
if match:
assert isinstance(module, Quantize)
assert module.signed == ref_values[
0], 'sign is not matched for {}'.format(name)
assert module.scale == ref_values[
1], 'scale is not matched for {}'.format(name)
def test_can_restore_binary_mask_on_magnitude_quant_algo_resume(tmp_path):
config = get_empty_config()
config["compression"] = [{
"algorithm": "magnitude_sparsity",
"params": {
"schedule": "multistep",
"multistep_sparsity_levels": [0.3, 0.5],
"weight_importance": "abs"
}
}, {
"algorithm": "quantization"
}]
sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
# load_state doesn't support CPU + Quantization
sparse_model = torch.nn.DataParallel(sparse_model)
sparse_model.cuda()
with torch.no_grad():
sparse_model(torch.ones([1, 1, 10, 10]))
config = get_empty_config()
config["compression"] = [{
"algorithm": "const_sparsity"
}, {
"algorithm": "quantization"
}]
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.get_nncf_wrapped_model().conv1.pre_ops['0']
check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.get_nncf_wrapped_model().conv2.pre_ops['0']
check_equal(ref_mask_2, op.operand.binary_mask)
def test_sparse_network(self, desc: ModelDesc, algo):
model = desc.model_builder()
from nncf.layers import NNCF_MODULES_MAP
sparsifiable_modules = list(NNCF_MODULES_MAP.values())
ref_num_sparsed = len(
get_all_modules_by_type(model, sparsifiable_modules))
config = get_empty_config(input_sample_sizes=desc.input_sample_sizes)
config["compression"] = {"algorithm": algo}
compressed_model, compression_ctrl = \
create_compressed_model_and_algo_for_test(model, config, dummy_forward_fn=desc.dummy_forward_fn)
assert ref_num_sparsed == len(compression_ctrl.sparsified_module_info)
check_model_graph(compressed_model, desc.dot_filename, algo)
def test_ordinary_load(algo, _model_wrapper, is_resume):
config = get_empty_config()
if algo:
config['compression'] = {'algorithm': algo, 'params': {}}
algo_save = create_test_compression_algo(config, BasicConvTestModel())
model_save = _model_wrapper['save_model'](algo_save.model)
algo_resume = create_test_compression_algo(config, BasicConvTestModel())
model_resume = _model_wrapper['resume_model'](algo_resume.model)
num_loaded = load_state(model_resume, model_save.state_dict(), is_resume)
assert num_loaded == len(model_save.state_dict())
def test_ordinary_load(algo, _model_wrapper, is_resume):
config = get_empty_config()
if algo:
config['compression'] = {'algorithm': algo}
compressed_model_save, _ = create_compressed_model_and_algo_for_test(BasicConvTestModel(), config)
model_save = _model_wrapper['save_model'](compressed_model_save)
compressed_model_resume, _ = create_compressed_model_and_algo_for_test(BasicConvTestModel(), config)
model_resume = _model_wrapper['resume_model'](compressed_model_resume)
num_loaded = load_state(model_resume, model_save.state_dict(), is_resume)
assert num_loaded == len(model_save.state_dict())
def test_ordinary_load(algo, _models, is_resume):
config = get_empty_config()
if algo:
config['compression'] = {'algorithm': algo, 'params': {}}
algo_save = create_test_compression_algo(config, _models['save_model'])
model_save = algo_save.model
algo_resume = create_test_compression_algo(config, _models['resume_model'])
model_resume = algo_resume.model
num_loaded = load_state(model_resume, model_save.state_dict(), is_resume)
assert num_loaded == len(model_save.state_dict())
def get_basic_quantization_config(qconfig, input_sample_sizes=None):
config = get_empty_config(input_sample_sizes=input_sample_sizes)
config['compression'] = {
'algorithm': 'quantization',
'activations': {
'mode': qconfig.mode,
'per_channel': qconfig.per_channel
},
'weights': {
'mode': qconfig.mode,
'per_channel': qconfig.per_channel
}
}
return config
def test_magnitude_scheduler_can_do_epoch_step__with_multistep():
_ = MagnitudeTestModel()
config = get_empty_config()
config["compression"] = {"algorithm": "magnitude_sparsity",
"params": {"schedule": "multistep", 'multistep_steps': [1]}}
_, compression_ctrl = create_compressed_model_and_algo_for_test(MagnitudeTestModel(), config)
scheduler = compression_ctrl.scheduler
assert isinstance(scheduler, MultiStepSparsityScheduler)
assert pytest.approx(compression_ctrl.sparsity_level) == 0.1
assert scheduler.sparsity_levels == [0.1, 0.5]
scheduler.epoch_step()
assert compression_ctrl.sparsity_level == 0.5
scheduler.epoch_step()
assert compression_ctrl.sparsity_level == 0.5
def test_activation_quantizers_order_is_the_same__for_resnet50(tmp_path):
config = get_empty_config(input_sample_size=[1, 3, 224, 224])
config['compression'] = {'algorithm': 'quantization'}
ngpus_per_node = torch.cuda.device_count()
torch.multiprocessing.spawn(activation_quantizers_dumping_worker,
nprocs=ngpus_per_node,
args=(config, tmp_path),
join=True)
with open(get_path_to_keys(tmp_path, 0), 'r') as f:
ref_list = f.readlines()
for i in range(1, ngpus_per_node):
with open(get_path_to_keys(tmp_path, i), 'r') as f:
curr_list = f.readlines()
assert curr_list == ref_list
def test_scheduler_can_do_epoch_step(self, algo, schedule, get_params, ref_levels):
model = BasicConvTestModel()
config = get_empty_config()
config['compression']['algorithm'] = algo
config['compression']["params"] = get_params()
config['compression']["params"]["schedule"] = schedule
compression_algo = create_compression_algorithm(model, config)
scheduler = compression_algo.scheduler
assert pytest.approx(scheduler.current_sparsity_level) == ref_levels[0]
for ref_level in ref_levels[1:]:
scheduler.epoch_step()
assert pytest.approx(scheduler.current_sparsity_level) == ref_level
for m in compression_algo.sparsified_module_info:
assert not m.operand.sparsify
def test_scale_and_sign_init_for_quant_algo__without_init_section(self, wrap_dataloader):
config = get_empty_config()
config['compression'] = {'algorithm': 'quantization'}
algo, compressed_model = self.create_algo_and_compressed_model(config)
device = next(compressed_model.parameters()).device
data_loader = self.create_dataloader(wrap_dataloader, config, device)
algo.initialize(data_loader)
self.check_sign_and_scale(compressed_model, {
'.*Sequential\\[0\\].*UpdateWeight.*': (True, 1),
'.*Sequential\\[1\\].*UpdateWeight. *': (False, 1),
'.*activation_quantizers.*Sequential\\[0\\].*': (True, 4),
'.*activation_quantizers.*Sequential\\[1\\].*': (True, 24)
})