def test_number_of_nodes_for_module_in_loop__not_input_node(self):
num_iter = 5
patch_torch_operators()
class LoopModule(nn.Module):
class Inner(nn.Module):
def forward(self, x):
s = F.sigmoid(x)
t = F.tanh(x)
result = F.sigmoid(x) * t + F.tanh(x) * s
return result
@staticmethod
def nodes_number():
return 7
def __init__(self):
super().__init__()
self.inner = self.Inner()
def forward(self, x):
for _ in range(num_iter):
x = self.inner(F.relu(x))
return x
def nodes_number(self):
return self.inner.nodes_number() + num_iter
test_module = LoopModule()
reset_context('test')
with context('test') as ctx:
_ = test_module(torch.zeros(1))
assert ctx.graph.get_nodes_count() == test_module.nodes_number()
def test_number_of_nodes_for_module_with_nested_loops(self):
num_iter = 5
patch_torch_operators()
class TestIterModule(nn.Module):
@ITERATION_MODULES.register()
class TestIterModule_ResetPoint(nn.Module):
def __init__(self, loop_module):
super().__init__()
self.loop_module = loop_module
def forward(self, x):
return self.loop_module(F.relu(x))
def __init__(self):
super().__init__()
self.loop_module = self.LoopModule2()
self.reset_point = self.TestIterModule_ResetPoint(
self.loop_module)
def forward(self, x):
for _ in range(num_iter):
x = self.reset_point(x)
return x
class LoopModule2(nn.Module):
@ITERATION_MODULES.register()
class LoopModule2_ResetPoint(nn.Module):
def __init__(self, inner):
super().__init__()
self.inner = inner
def forward(self, x):
return self.inner(F.relu(x))
def __init__(self):
super().__init__()
self.inner = self.Inner()
self.reset_helper = self.LoopModule2_ResetPoint(self.inner)
def forward(self, x):
for _ in range(num_iter):
self.reset_helper(x)
return x
class Inner(nn.Module):
def forward(self, x):
s = F.sigmoid(x)
t = F.tanh(x)
result = t + s
return result
test_module = TestIterModule()
reset_context('test')
with context('test') as ctx:
_ = test_module(torch.zeros(1))
assert ctx.graph.get_nodes_count() == num_iter
def test_build_graph(self, model_name, model_builder, input_size):
net = model_builder()
ctx = reset_context('test')
with context('test') as c:
_ = net(torch.zeros(input_size))
c.reset_scope_operator_call_counters()
_ = net(torch.zeros(input_size))
check_graph(to_networkx(ctx), model_name, 'original')
def test_build_graph(self, model_name, model_builder, forward_fn_):
net = model_builder()
net.to(self.device)
ctx = reset_context('test')
with context('test') as c:
forward_fn_(net)
c.reset_scope_operator_call_counters()
forward_fn_(net)
check_graph(ctx.graph, model_name, 'original')
def build_graph(self, model: torch.nn.Module, context_name: str) -> NNCFGraph:
logger.info("Building graph with context: {}".format(context_name))
sd = deepcopy(model.state_dict())
ctx = get_context(context_name)
with context(context_name):
self.custom_forward_fn(model)
model.load_state_dict(sd)
if isinstance(model, PostGraphBuildActing):
model.post_build_graph_actions()
return ctx.graph
def test_iterate_module_list():
class Net(nn.Module):
def __init__(self):
super().__init__()
self.ml = nn.ModuleList([nn.Conv2d(1, 1, 1), nn.Conv2d(1, 1, 1)])
def forward(self, x):
return [self.ml[0](x), self.ml[1](x)]
net = Net()
ctx = reset_context('test')
with context('test'):
_ = net(torch.zeros(1, 1, 1, 1))
check_graph(ctx.graph, 'case_iterate_module_list.dot', 'original')
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))
ctx = reset_context('test')
config = get_empty_config(input_sample_size=input_size)
config["compression"] = {"algorithm": algo, "params": params}
compression_algo = create_compression_algorithm(model, config)
assert ref_num_sparsed == len(compression_algo.sparsified_module_info)
model = compression_algo.model
with context('test') as c:
_ = model(torch.zeros(input_size))
c.reset_scope_operator_call_counters()
_ = model(torch.zeros(input_size))
check_graph(to_networkx(ctx), model_name, algo)
def test_sparse_network(self, model_name, model_builder, forward_fn_, 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))
ctx = reset_context('test')
config = get_empty_config()
config["compression"] = {"algorithm": algo, "params": params}
compression_algo = create_compression_algorithm(model, config, dummy_forward_fn=forward_fn_)
assert ref_num_sparsed == len(compression_algo.sparsified_module_info)
model = compression_algo.model
model.to(self.device)
with context('test') as c:
forward_fn_(model)
c.reset_scope_operator_call_counters()
forward_fn_(model)
check_graph(ctx.graph, model_name, algo)
def test_number_of_nodes_for_module_in_loop(self):
num_iter = 5
patch_torch_operators()
class LoopModule(nn.Module):
class Inner(nn.Module):
def __init__(self):
super().__init__()
self.operator1 = torch.sigmoid
self.operator2 = torch.tanh
def forward(self, x):
s = self.operator1(x)
t = self.operator2(x)
result = t + s
return result
@staticmethod
def nodes_number():
return 3
def __init__(self):
super().__init__()
self.inner = self.Inner()
def forward(self, x):
for _ in range(num_iter):
x = self.inner(x)
return x
def nodes_number(self):
return self.inner.nodes_number()
test_module = LoopModule()
reset_context('test')
with context('test') as ctx:
_ = test_module(torch.zeros(1))
assert ctx.graph.get_nodes_count() == test_module.nodes_number()
def test_number_of_nodes_for_repeated_module(self):
patch_torch_operators()
class LoopModule(nn.Module):
def __init__(self):
super().__init__()
self.operator = F.relu
self.layers = nn.ModuleList(
[nn.Conv2d(1, 1, 1),
nn.Conv2d(1, 1, 1)])
def forward(self, x):
for layer in self.layers:
x = F.relu(layer(x))
return x
test_module = LoopModule()
reset_context('test')
with context('test') as ctx:
x = test_module(torch.zeros(1, 1, 1, 1))
assert ctx.graph.get_nodes_count() == 4
_ = test_module(x)
assert ctx.graph.get_nodes_count() == 8
def test_number_of_calling_fq_for_gnmt(self, tmp_path):
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)
patch_torch_operators()
config = get_empty_config(input_sample_size=input_sample_size)
config['compression'] = \
{'algorithm': 'quantization',
'quantize_inputs': True,
'quantizable_subgraph_patterns': [["linear", "__add__"],
["sigmoid", "__mul__", "__add__"],
["__add__", "tanh", "__mul__"],
["sigmoid", "__mul__"]],
'scopes_without_shape_matching':
['GNMT/ResidualRecurrentDecoder[decoder]/RecurrentAttention[att_rnn]/BahdanauAttention[attn]'],
'disable_function_quantization_hooks': True}
config.log_dir = str(tmp_path)
reset_context('orig')
reset_context('quantized_graphs')
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.forward(input_encoder, input_enc_len, input_decoder)
_, model = create_compressed_model(model, config, dummy_forward_fn)
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 model.all_quantizations.items():
counter = Counter()
counters[name] = counter
quantizer.register_forward_pre_hook(partial(hook, counter=counter))
with context('quantized_graphs') as ctx:
dummy_forward_fn(model)
assert ctx.graph.get_nodes_count() == 239
assert len(counters) == 68
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 ctx.graph.get_nodes_count() == 239
assert len(counters) == 68
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
def test_number_of_calling_fq_for_lstm(self, tmp_path):
p = LSTMTestSizes(1, 1, 1, 5)
num_layers = 2
bidirectional = True
num_directions = 2 if bidirectional else 1
bias = True
batch_first = False
patch_torch_operators()
config = get_empty_config(input_sample_size=(p.seq_length, p.batch,
p.input_size))
config['compression'] = {
'algorithm': 'quantization',
'quantize_inputs': True
}
config.log_dir = str(tmp_path)
reset_context('orig')
reset_context('quantized_graphs')
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)
_ = reset_context('orig')
_ = reset_context('quantized_graphs')
_, model = create_compressed_model(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 model.all_quantizations.items():
counter = Counter()
counters[name] = counter
quantizer.register_forward_pre_hook(partial(hook, counter=counter))
with context('quantized_graphs') as ctx:
_ = model(test_data.x, test_hidden)
assert ctx.graph.get_nodes_count() == 110
ctx.graph.dump_graph(
os.path.join(config.log_dir, "compressed_graph_next.dot"))
assert len(counters) == 54
for counter in counters.values():
assert counter.count == p.seq_length
