def test_sparsenn_shadow(self):
sparse_nn = SparseNNModel().eval()
# quantize the embeddings and the dense part separately, using FX graph mode
sparse_nn.dense_top = prepare_fx(
sparse_nn.dense_top,
{'': torch.quantization.default_qconfig},
)
# calibrate
idx = torch.LongTensor([1, 2, 4, 5, 4, 3, 2, 9])
offsets = torch.LongTensor([0, 4])
x = torch.randn(2, 4)
sparse_nn(idx, offsets, x)
# convert
sparse_nn_q = copy.deepcopy(sparse_nn)
sparse_nn_q.dense_top = convert_fx(sparse_nn_q.dense_top)
# test out compare shadow activations API
sparse_nn_q.dense_top = prepare_model_with_stubs(
'fp32_prepared', sparse_nn.dense_top,
'int8', sparse_nn_q.dense_top, OutputLogger)
# calibrate
sparse_nn_q(idx, offsets, x)
# check activation result correctness
act_compare_dict = get_matching_activations_a_shadows_b(
sparse_nn_q, OutputLogger)
self.assertTrue(len(act_compare_dict) == 4)
self.assert_ns_compare_dict_valid(act_compare_dict)
def test_sparsenn_shadow(self):
for should_log_inputs in (True, False):
sparse_nn = SparseNNModel().eval()
idx = torch.LongTensor([1, 2, 4, 5, 4, 3, 2, 9])
offsets = torch.LongTensor([0, 4])
x = torch.randn(2, 4)
self._test_match_activations(sparse_nn, (idx, offsets, x),
results_len=4,
should_log_inputs=should_log_inputs)