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_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']
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_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 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_sizes=[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 = {}
counter_for_input_quantizer = None
for name, quantizer in algo.all_quantizations.items():
counter = Counter()
quantizer.register_forward_pre_hook(partial(hook, counter=counter))
if str(name) == '/nncf_model_input_0|OUTPUT':
counter_for_input_quantizer = counter
continue
counters[name] = counter
_ = model(test_data.x, test_hidden)
assert model.get_graph().get_nodes_count(
) == 112 # NB: may always fail in debug due to superfluous 'cat' nodes
assert len(counters) + 1 == 55 # 8 WQ + 46 AQ + 1 input AQ
for counter in counters.values():
assert counter.count == p.seq_length
assert counter_for_input_quantizer.count == 1
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",
"params": {"sparsity_init": 0.1,
"sparsity_target": 0.9,
"sparsity_target_epoch": 3,
"sparsity_freeze_epoch": 2}}
_, 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, 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_scheduler_can_do_epoch_step(self, algo, schedule, get_params,
ref_levels):
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {
'algorithm': algo,
"sparsity_init": 0.2,
"params": {
**get_params(), "schedule": schedule
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
scheduler = compression_ctrl.scheduler
scheduler.epoch_step()
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_ctrl.sparsified_module_info:
if hasattr(m.operand, "frozen"):
assert m.operand.frozen
def test_create_rb_algo_with_per_layer_loss():
config = get_empty_config()
config['compression'] = {'algorithm': 'rb_sparsity', "params": {"sparsity_level_setting_mode": 'local'}}
_, compression_ctrl = create_compressed_model_and_algo_for_test(MockModel(), config)
# pylint: disable=protected-access
assert isinstance(compression_ctrl._loss, SparseLossForPerLayerSparsity)
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'}
# 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 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["target_device"] = hw_config_type
config["compression"] = {
"algorithm": "quantization",
}
return config
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_rb_sparsity__can_set_sparsity_level_for_module():
config = get_empty_config()
config['compression'] = {'algorithm': 'rb_sparsity', "params": {"sparsity_level_setting_mode": 'local'}}
_, compression_ctrl = create_compressed_model_and_algo_for_test(MockModel(), config)
# pylint: disable=protected-access
assert list(compression_ctrl._loss.per_layer_target.values())[0] == 1
compression_ctrl.set_sparsity_level(0.7, compression_ctrl.sparsified_module_info[0])
assert list(compression_ctrl._loss.per_layer_target.values())[0] == pytest.approx(0.3)
def test_can_create_rb_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression'] = {
'algorithm': 'rb_sparsity',
"params": {
"schedule": 'adaptive'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, AdaptiveSparsityScheduler)
def create_config():
config = get_empty_config()
config['compression'] = {
'algorithm': 'quantization',
'initializer': {
'range': {
'num_init_steps': 1
}
}
}
return config
def test_can_choose_scheduler(algo, schedule_type, scheduler_class):
config = get_empty_config()
config['compression'] = {
'algorithm': algo,
"params": {
"schedule": schedule_type
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, scheduler_class)
def get_basic_quantization_config(quantization_type, input_sample_sizes=None):
config = get_empty_config(input_sample_sizes=input_sample_sizes)
config["compression"] = {
"algorithm": "quantization",
"activations": {
"mode": quantization_type
},
"weights": {
"mode": quantization_type
}
}
return config
def test_can_freeze_binary_masks():
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {'algorithm': "magnitude_sparsity"}
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
for sparse_layer in compression_ctrl.sparsified_module_info:
assert not sparse_layer.operand.frozen
compression_ctrl.freeze()
for sparse_layer in compression_ctrl.sparsified_module_info:
assert sparse_layer.operand.frozen
def test_create_rb_algo_with_stub_scheduler():
config = get_empty_config()
config['compression'] = {
'algorithm': 'rb_sparsity',
"params": {
"sparsity_level_setting_mode": 'local'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
# pylint: disable=protected-access
assert isinstance(compression_ctrl.scheduler, StubCompressionScheduler)
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_activation_quantizers_order_is_the_same__for_resnet50(tmp_path):
config = get_empty_config(input_sample_sizes=[1, 3, 224, 224])
config['compression'] = {'algorithm': 'quantization', "initializer": {"range": {"num_init_steps": 0}}}
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_can_create_sparse_scheduler__with_defaults(self, algo):
config = get_empty_config()
config['compression'] = {
'algorithm': algo,
"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.sparsity_target == 0.5
assert scheduler.sparsity_target_epoch == 90
assert scheduler.sparsity_freeze_epoch == 100
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_export_lstm_cell(tmp_path):
config = get_empty_config(model_size=1, input_sample_sizes=[1, 1])
config['compression'] = {'algorithm': 'quantization'}
model, algo = create_compressed_model_and_algo_for_test(
LSTMCellNNCF(1, 1), 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 == 12
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(scheduler.current_sparsity_level) == 0.1
assert scheduler.sparsity_levels == [0.1, 0.5]
scheduler.epoch_step()
assert scheduler.current_sparsity_level == 0.5
scheduler.epoch_step()
assert scheduler.current_sparsity_level == 0.5
def test_can_create_const_sparse_algo__with_default():
model = BasicConvTestModel()
config = get_empty_config()
config["compression"] = {"algorithm": "const_sparsity"}
sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(
deepcopy(model), config)
assert isinstance(compression_ctrl, ConstSparsityController)
assert len(list(sparse_model.modules())) == 7
_, sparse_model_conv = check_correct_nncf_modules_replacement(
model, sparse_model)
for sparse_module in sparse_model_conv.values():
store = []
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 create_empty_config_without_init_section():
config = get_empty_config()
config['compression'] = {'algorithm': 'quantization'}
return config
def get_const_sparsity_config():
config = get_empty_config()
config['compression'] = {'algorithm': 'const_sparsity'}
return config
def create_regular_dataloader():
return create_mock_dataloader(config=get_empty_config(),
num_samples=N_SAMPLE)
def __init__(self):
self._config = get_empty_config()
self._algorithm_sections = {}
def test_number_of_calling_fq_for_gnmt(self):
torch.cuda.set_device(0)
device = torch.device('cuda')
batch_first = False
vocab_size = 32000
model_config = {
'hidden_size': 100,
'vocab_size': vocab_size,
'num_layers': 4,
'dropout': 0.2,
'batch_first': batch_first,
'share_embedding': True,
}
batch_size = 128
sequence_size = 50
input_sample_size = [batch_size, sequence_size
] if batch_first else [sequence_size, batch_size]
config = get_empty_config(input_sample_sizes=input_sample_size)
config['compression'] = \
{'algorithm': 'quantization',
'quantize_inputs': True,
'quantizable_subgraph_patterns': [["linear", "__add__"],
["sigmoid", "__mul__", "__add__"],
["__add__", "tanh", "__mul__"],
["sigmoid", "__mul__"]],
'disable_function_quantization_hooks': True}
config['scopes_without_shape_matching'] = \
['GNMT/ResidualRecurrentDecoder[decoder]/RecurrentAttention[att_rnn]/BahdanauAttention[attn]', ]
model = GNMT(**model_config)
model = replace_lstm(model)
model.to(device)
def dummy_forward_fn(model, seq_len=sequence_size):
def gen_packed_sequence():
seq_list = []
seq_lens = torch.LongTensor(batch_size).random_(1, seq_len + 1)
seq_lens = torch.sort(seq_lens, descending=True).values
for seq_size in seq_lens:
seq_list.append(
torch.LongTensor(seq_size.item()).random_(
1, vocab_size).to(device))
padded_seq_batch = torch.nn.utils.rnn.pad_sequence(
seq_list, batch_first=batch_first)
return padded_seq_batch, seq_lens
x_data, seq_lens = gen_packed_sequence()
input_encoder = x_data
input_enc_len = seq_lens.to(device)
input_decoder = gen_packed_sequence()[0]
model(input_encoder, input_enc_len, input_decoder)
algo, model = create_compressed_model(
model,
config,
dummy_forward_fn=dummy_forward_fn,
dump_graphs=False)
model.to(device)
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[str(name)] = counter
quantizer.register_forward_pre_hook(partial(hook, counter=counter))
dummy_forward_fn(model)
assert model.get_graph().get_nodes_count(
) == 232 # NB: may always fail in debug due to superfluous 'cat' nodes
assert len(counters) == 57
for name, counter in counters.items():
if 'cell' in name or "LSTMCellForwardNNCF" in name:
assert counter.count == sequence_size, name
else:
assert counter.count == 1, name
new_seq_len = int(sequence_size / 2)
dummy_forward_fn(model, new_seq_len)
assert model.get_graph().get_nodes_count(
) == 232 # NB: may always fail in debug due to superfluous 'cat' nodes
assert len(counters) == 57
for name, counter in counters.items():
if 'cell' in name or "LSTMCellForwardNNCF" in name:
assert counter.count == sequence_size + new_seq_len, name
else:
assert counter.count == 2, name